bartz 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

bartz/mcmcstep.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/mcmcstep.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -37,14 +37,14 @@ import functools
 import math
 
 import jax
-from jax import random
+from jax import lax, random
 from jax import numpy as jnp
-from jax import lax
 
-from . import jaxext
-from . import grove
+from . import grove, jaxext
 
-def init(*,
+
+def init(
+    *,
     X,
     y,
     max_split,
@@ -52,13 +52,14 @@ def init(*,
     p_nonterminal,
     sigma2_alpha,
     sigma2_beta,
+    error_scale=None,
     small_float=jnp.float32,
     large_float=jnp.float32,
     min_points_per_leaf=None,
     resid_batch_size='auto',
     count_batch_size='auto',
     save_ratios=False,
-    ):
+):
     """
     Make a BART posterior sampling MCMC initial state.
 
@@ -76,9 +77,12 @@ def init(*,
         The probability of a nonterminal node at each depth. The maximum depth
         of trees is fixed by the length of this array.
     sigma2_alpha : float
-        The shape parameter of the inverse gamma prior on the noise variance.
+        The shape parameter of the inverse gamma prior on the error variance.
     sigma2_beta : float
-        The scale parameter of the inverse gamma prior on the noise variance.
+        The scale parameter of the inverse gamma prior on the error variance.
+    error_scale : float array (n,), optional
+        Each error is scaled by the corresponding factor in `error_scale`, so
+        the error variance for ``y[i]`` is ``sigma2 * error_scale[i] ** 2``.
     small_float : dtype, default float32
         The dtype for large arrays used in the algorithm.
     large_float : dtype, default float32
@@ -110,6 +114,8 @@ def init(*,
             roundoff.
         'sigma2' : large_float
             The noise variance.
+        'prec_scale' : large_float array (n,) or None
+            The scale on the error precision, i.e., ``1 / error_scale ** 2``.
         'grow_prop_count', 'prune_prop_count' : int
             The number of grow/prune proposals made during one full MCMC cycle.
         'grow_acc_count', 'prune_acc_count' : int
@@ -169,16 +175,27 @@ def init(*,
     small_float = jnp.dtype(small_float)
     large_float = jnp.dtype(large_float)
     y = jnp.asarray(y, small_float)
-    resid_batch_size, count_batch_size = _choose_suffstat_batch_size(resid_batch_size, count_batch_size, y, 2 ** max_depth * num_trees)
+    resid_batch_size, count_batch_size = _choose_suffstat_batch_size(
+        resid_batch_size, count_batch_size, y, 2**max_depth * num_trees
+    )
     sigma2 = jnp.array(sigma2_beta / sigma2_alpha, large_float)
-    sigma2 = jnp.where(jnp.isfinite(sigma2) & (sigma2 > 0), sigma2, 1)
+    sigma2 = jnp.where(
+        jnp.isfinite(sigma2) & (sigma2 > 0), sigma2, 1
+    )  # TODO: I don't like this error check, these functions should be low-level and just do the thing. Why is it here?
 
     bart = dict(
         leaf_trees=make_forest(max_depth, small_float),
-        var_trees=make_forest(max_depth - 1, jaxext.minimal_unsigned_dtype(X.shape[0] - 1)),
+        var_trees=make_forest(
+            max_depth - 1, jaxext.minimal_unsigned_dtype(X.shape[0] - 1)
+        ),
         split_trees=make_forest(max_depth - 1, max_split.dtype),
         resid=jnp.asarray(y, large_float),
         sigma2=sigma2,
+        prec_scale=(
+            None
+            if error_scale is None
+            else lax.reciprocal(jnp.square(jnp.asarray(error_scale, large_float)))
+        ),
         grow_prop_count=jnp.zeros((), int),
         grow_acc_count=jnp.zeros((), int),
         prune_prop_count=jnp.zeros((), int),
@@ -190,14 +207,18 @@ def init(*,
         max_split=jnp.asarray(max_split),
         y=y,
         X=jnp.asarray(X),
-        leaf_indices=jnp.ones((num_trees, y.size), jaxext.minimal_unsigned_dtype(2 ** max_depth - 1)),
+        leaf_indices=jnp.ones(
+            (num_trees, y.size), jaxext.minimal_unsigned_dtype(2**max_depth - 1)
+        ),
         min_points_per_leaf=(
-            None if min_points_per_leaf is None else
-            jnp.asarray(min_points_per_leaf)
+            None if min_points_per_leaf is None else jnp.asarray(min_points_per_leaf)
         ),
         affluence_trees=(
-            None if min_points_per_leaf is None else
-            make_forest(max_depth - 1, bool).at[:, 1].set(y.size >= 2 * min_points_per_leaf)
+            None
+            if min_points_per_leaf is None
+            else make_forest(max_depth - 1, bool)
+            .at[:, 1]
+            .set(y.size >= 2 * min_points_per_leaf)
         ),
         opt=jaxext.LeafDict(
             small_float=small_float,
@@ -216,8 +237,8 @@ def init(*,
 
     return bart
 
-def _choose_suffstat_batch_size(resid_batch_size, count_batch_size, y, forest_size):
 
+def _choose_suffstat_batch_size(resid_batch_size, count_batch_size, y, forest_size):
     @functools.cache
     def get_platform():
         try:
@@ -233,9 +254,9 @@ def _choose_suffstat_batch_size(resid_batch_size, count_batch_size, y, forest_si
         platform = get_platform()
         n = max(1, y.size)
         if platform == 'cpu':
-            resid_batch_size = 2 ** int(round(math.log2(n / 6))) # n/6
+            resid_batch_size = 2 ** int(round(math.log2(n / 6)))  # n/6
         elif platform == 'gpu':
-            resid_batch_size = 2 ** int(round((1 + math.log2(n)) / 3)) # n^1/3
+            resid_batch_size = 2 ** int(round((1 + math.log2(n)) / 3))  # n^1/3
         resid_batch_size = max(1, resid_batch_size)
 
     if count_batch_size == 'auto':
@@ -244,9 +265,9 @@ def _choose_suffstat_batch_size(resid_batch_size, count_batch_size, y, forest_si
             count_batch_size = None
         elif platform == 'gpu':
             n = max(1, y.size)
-            count_batch_size = 2 ** int(round(math.log2(n) / 2 - 2)) # n^1/2
-                # /4 is good on V100, /2 on L4/T4, still haven't tried A100
-            max_memory = 2 ** 29
+            count_batch_size = 2 ** int(round(math.log2(n) / 2 - 2))  # n^1/2
+            # /4 is good on V100, /2 on L4/T4, still haven't tried A100
+            max_memory = 2**29
             itemsize = 4
             min_batch_size = int(math.ceil(forest_size * itemsize * n / max_memory))
             count_batch_size = max(count_batch_size, min_batch_size)
@@ -254,16 +275,17 @@ def _choose_suffstat_batch_size(resid_batch_size, count_batch_size, y, forest_si
 
     return resid_batch_size, count_batch_size
 
-def step(bart, key):
+
+def step(key, bart):
     """
     Perform one full MCMC step on a BART state.
 
     Parameters
     ----------
-    bart : dict
-        A BART mcmc state, as created by `init`.
     key : jax.dtypes.prng_key array
         A jax random key.
+    bart : dict
+        A BART mcmc state, as created by `init`.
 
     Returns
     -------
@@ -271,19 +293,20 @@ def step(bart, key):
         The new BART mcmc state.
     """
     key, subkey = random.split(key)
-    bart = sample_trees(bart, subkey)
-    return sample_sigma(bart, key)
+    bart = sample_trees(subkey, bart)
+    return sample_sigma(key, bart)
+
 
-def sample_trees(bart, key):
+def sample_trees(key, bart):
     """
     Forest sampling step of BART MCMC.
 
     Parameters
     ----------
-    bart : dict
-        A BART mcmc state, as created by `init`.
     key : jax.dtypes.prng_key array
         A jax random key.
+    bart : dict
+        A BART mcmc state, as created by `init`.
 
     Returns
     -------
@@ -295,19 +318,20 @@ def sample_trees(bart, key):
     This function zeroes the proposal counters.
     """
     key, subkey = random.split(key)
-    moves = sample_moves(bart, subkey)
-    return accept_moves_and_sample_leaves(bart, moves, key)
+    moves = sample_moves(subkey, bart)
+    return accept_moves_and_sample_leaves(key, bart, moves)
 
-def sample_moves(bart, key):
+
+def sample_moves(key, bart):
     """
     Propose moves for all the trees.
 
     Parameters
     ----------
-    bart : dict
-        BART mcmc state.
     key : jax.dtypes.prng_key array
         A jax random key.
+    bart : dict
+        BART mcmc state.
 
     Returns
     -------
@@ -343,14 +367,22 @@ def sample_moves(bart, key):
     key, subkey = key[0], key[1:]
 
     # compute moves
-    grow_moves, prune_moves = _sample_moves_vmap_trees(bart['var_trees'], bart['split_trees'], bart['affluence_trees'], bart['max_split'], bart['p_nonterminal'], bart['p_propose_grow'], subkey)
+    grow_moves, prune_moves = _sample_moves_vmap_trees(
+        subkey,
+        bart['var_trees'],
+        bart['split_trees'],
+        bart['affluence_trees'],
+        bart['max_split'],
+        bart['p_nonterminal'],
+        bart['p_propose_grow'],
+    )
 
     u, logu = random.uniform(key, (2, ntree), bart['opt']['large_float'])
 
     # choose between grow or prune
     grow_allowed = grow_moves['num_growable'].astype(bool)
     p_grow = jnp.where(grow_allowed & prune_moves['allowed'], 0.5, grow_allowed)
-    grow = u < p_grow # use < instead of <= because u is in [0, 1)
+    grow = u < p_grow  # use < instead of <= because u is in [0, 1)
 
     # compute children indices
     node = jnp.where(grow, grow_moves['node'], prune_moves['node'])
@@ -364,22 +396,28 @@ def sample_moves(bart, key):
         node=node,
         left=left,
         right=right,
-        partial_ratio=jnp.where(grow, grow_moves['partial_ratio'], prune_moves['partial_ratio']),
+        partial_ratio=jnp.where(
+            grow, grow_moves['partial_ratio'], prune_moves['partial_ratio']
+        ),
         grow_var=grow_moves['var'],
         grow_split=grow_moves['split'],
         var_trees=grow_moves['var_tree'],
         logu=jnp.log1p(-logu),
     )
 
-@functools.partial(jaxext.vmap_nodoc, in_axes=(0, 0, 0, None, None, None, 0))
+
+@functools.partial(jaxext.vmap_nodoc, in_axes=(0, 0, 0, 0, None, None, None))
 def _sample_moves_vmap_trees(*args):
-    args, key = args[:-1], args[-1]
+    key, args = args[0], args[1:]
     key, key1 = random.split(key)
-    grow = grow_move(*args, key)
-    prune = prune_move(*args, key1)
+    grow = grow_move(key, *args)
+    prune = prune_move(key1, *args)
     return grow, prune
 
-def grow_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, p_propose_grow, key):
+
+def grow_move(
+    key, var_tree, split_tree, affluence_tree, max_split, p_nonterminal, p_propose_grow
+):
     """
     Tree structure grow move proposal of BART MCMC.
 
@@ -426,14 +464,18 @@ def grow_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, p_
 
     key, key1, key2 = random.split(key, 3)
 
-    leaf_to_grow, num_growable, prob_choose, num_prunable = choose_leaf(split_tree, affluence_tree, p_propose_grow, key)
+    leaf_to_grow, num_growable, prob_choose, num_prunable = choose_leaf(
+        key, split_tree, affluence_tree, p_propose_grow
+    )
 
-    var = choose_variable(var_tree, split_tree, max_split, leaf_to_grow, key1)
+    var = choose_variable(key1, var_tree, split_tree, max_split, leaf_to_grow)
     var_tree = var_tree.at[leaf_to_grow].set(var.astype(var_tree.dtype))
 
-    split = choose_split(var_tree, split_tree, max_split, leaf_to_grow, key2)
+    split = choose_split(key2, var_tree, split_tree, max_split, leaf_to_grow)
 
-    ratio = compute_partial_ratio(prob_choose, num_prunable, p_nonterminal, leaf_to_grow)
+    ratio = compute_partial_ratio(
+        prob_choose, num_prunable, p_nonterminal, leaf_to_grow
+    )
 
     return dict(
         num_growable=num_growable,
@@ -444,7 +486,8 @@ def grow_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, p_
         var_tree=var_tree,
     )
 
-def choose_leaf(split_tree, affluence_tree, p_propose_grow, key):
+
+def choose_leaf(key, split_tree, affluence_tree, p_propose_grow):
     """
     Choose a leaf node to grow in a tree.
 
@@ -482,6 +525,7 @@ def choose_leaf(split_tree, affluence_tree, p_propose_grow, key):
     num_prunable = jnp.count_nonzero(is_parent)
     return leaf_to_grow, num_growable, prob_choose, num_prunable
 
+
 def growable_leaves(split_tree, affluence_tree):
     """
     Return a mask indicating the leaf nodes that can be proposed for growth.
@@ -505,6 +549,7 @@ def growable_leaves(split_tree, affluence_tree):
         is_growable &= affluence_tree
     return is_growable
 
+
 def categorical(key, distr):
     """
     Return a random integer from an arbitrary distribution.
@@ -521,12 +566,15 @@ def categorical(key, distr):
     u : int
         A random integer in the range ``[0, n)``. If all probabilities are zero,
         return ``n``.
+    norm : float
+        The sum of `distr`.
     """
     ecdf = jnp.cumsum(distr)
     u = random.uniform(key, (), ecdf.dtype, 0, ecdf[-1])
     return jnp.searchsorted(ecdf, u, 'right'), ecdf[-1]
 
-def choose_variable(var_tree, split_tree, max_split, leaf_index, key):
+
+def choose_variable(key, var_tree, split_tree, max_split, leaf_index):
     """
     Choose a variable to split on for a new non-terminal node.
 
@@ -556,6 +604,7 @@ def choose_variable(var_tree, split_tree, max_split, leaf_index, key):
     var_to_ignore = fully_used_variables(var_tree, split_tree, max_split, leaf_index)
     return randint_exclude(key, max_split.size, var_to_ignore)
 
+
 def fully_used_variables(var_tree, split_tree, max_split, leaf_index):
     """
     Return a list of variables that have an empty split range at a given node.
@@ -586,6 +635,7 @@ def fully_used_variables(var_tree, split_tree, max_split, leaf_index):
     num_split = r - l
     return jnp.where(num_split == 0, var_to_ignore, max_split.size)
 
+
 def ancestor_variables(var_tree, max_split, node_index):
     """
     Return the list of variables in the ancestors of a node.
@@ -606,8 +656,11 @@ def ancestor_variables(var_tree, max_split, node_index):
         the parent. Unused spots are filled with `p`.
     """
     max_num_ancestors = grove.tree_depth(var_tree) - 1
-    ancestor_vars = jnp.zeros(max_num_ancestors, jaxext.minimal_unsigned_dtype(max_split.size))
+    ancestor_vars = jnp.zeros(
+        max_num_ancestors, jaxext.minimal_unsigned_dtype(max_split.size)
+    )
     carry = ancestor_vars.size - 1, node_index, ancestor_vars
+
     def loop(carry, _):
         i, index, ancestor_vars = carry
         index >>= 1
@@ -615,9 +668,11 @@ def ancestor_variables(var_tree, max_split, node_index):
         var = jnp.where(index, var, max_split.size)
         ancestor_vars = ancestor_vars.at[i].set(var)
         return (i - 1, index, ancestor_vars), None
+
     (_, _, ancestor_vars), _ = lax.scan(loop, carry, None, ancestor_vars.size)
     return ancestor_vars
 
+
 def split_range(var_tree, split_tree, max_split, node_index, ref_var):
     """
     Return the range of allowed splits for a variable at a given node.
@@ -641,8 +696,11 @@ def split_range(var_tree, split_tree, max_split, node_index, ref_var):
         The range of allowed splits is [l, r).
     """
     max_num_ancestors = grove.tree_depth(var_tree) - 1
-    initial_r = 1 + max_split.at[ref_var].get(mode='fill', fill_value=0).astype(jnp.int32)
+    initial_r = 1 + max_split.at[ref_var].get(mode='fill', fill_value=0).astype(
+        jnp.int32
+    )
     carry = 0, initial_r, node_index
+
     def loop(carry, _):
         l, r, index = carry
         right_child = (index & 1).astype(bool)
@@ -652,9 +710,11 @@ def split_range(var_tree, split_tree, max_split, node_index, ref_var):
         l = jnp.where(cond & right_child, jnp.maximum(l, split), l)
         r = jnp.where(cond & ~right_child, jnp.minimum(r, split), r)
         return (l, r, index), None
+
     (l, r, _), _ = lax.scan(loop, carry, None, max_num_ancestors)
     return l + 1, r
 
+
 def randint_exclude(key, sup, exclude):
     """
     Return a random integer in a range, excluding some values.
@@ -679,12 +739,15 @@ def randint_exclude(key, sup, exclude):
     exclude = jnp.unique(exclude, size=exclude.size, fill_value=sup)
     num_allowed = sup - jnp.count_nonzero(exclude < sup)
     u = random.randint(key, (), 0, num_allowed)
+
     def loop(u, i):
         return jnp.where(i <= u, u + 1, u), None
+
     u, _ = lax.scan(loop, u, exclude)
     return u
 
-def choose_split(var_tree, split_tree, max_split, leaf_index, key):
+
+def choose_split(key, var_tree, split_tree, max_split, leaf_index):
     """
     Choose a split point for a new non-terminal node.
 
@@ -711,6 +774,7 @@ def choose_split(var_tree, split_tree, max_split, leaf_index, key):
     l, r = split_range(var_tree, split_tree, max_split, leaf_index, var)
     return random.randint(key, (), l, r)
 
+
 def compute_partial_ratio(prob_choose, num_prunable, p_nonterminal, leaf_to_grow):
     """
     Compute the product of the transition and prior ratios of a grow move.
@@ -742,9 +806,9 @@ def compute_partial_ratio(prob_choose, num_prunable, p_nonterminal, leaf_to_grow
     # computed in the acceptance phase
 
     prune_allowed = leaf_to_grow != 1
-        # prune allowed  <--->  the initial tree is not a root
-        # leaf to grow is root  -->  the tree can only be a root
-        # tree is a root  -->  the only leaf I can grow is root
+    # prune allowed  <--->  the initial tree is not a root
+    # leaf to grow is root  -->  the tree can only be a root
+    # tree is a root  -->  the only leaf I can grow is root
 
     p_grow = jnp.where(prune_allowed, 0.5, 1)
 
@@ -757,7 +821,10 @@ def compute_partial_ratio(prob_choose, num_prunable, p_nonterminal, leaf_to_grow
 
     return tree_ratio / inv_trans_ratio
 
-def prune_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, p_propose_grow, key):
+
+def prune_move(
+    key, var_tree, split_tree, affluence_tree, max_split, p_nonterminal, p_propose_grow
+):
     """
     Tree structure prune move proposal of BART MCMC.
 
@@ -792,18 +859,23 @@ def prune_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, p
             the likelihood ratio and the probability of proposing the prune
             move. This ratio is inverted.
     """
-    node_to_prune, num_prunable, prob_choose = choose_leaf_parent(split_tree, affluence_tree, p_propose_grow, key)
-    allowed = split_tree[1].astype(bool) # allowed iff the tree is not a root
+    node_to_prune, num_prunable, prob_choose = choose_leaf_parent(
+        key, split_tree, affluence_tree, p_propose_grow
+    )
+    allowed = split_tree[1].astype(bool)  # allowed iff the tree is not a root
 
-    ratio = compute_partial_ratio(prob_choose, num_prunable, p_nonterminal, node_to_prune)
+    ratio = compute_partial_ratio(
+        prob_choose, num_prunable, p_nonterminal, node_to_prune
+    )
 
     return dict(
         allowed=allowed,
         node=node_to_prune,
-        partial_ratio=ratio, # it is inverted in accept_move_and_sample_leaves
+        partial_ratio=ratio,  # it is inverted in accept_move_and_sample_leaves
     )
 
-def choose_leaf_parent(split_tree, affluence_tree, p_propose_grow, key):
+
+def choose_leaf_parent(key, split_tree, affluence_tree, p_propose_grow):
     """
     Pick a non-terminal node with leaf children to prune in a tree.
 
@@ -835,8 +907,7 @@ def choose_leaf_parent(split_tree, affluence_tree, p_propose_grow, key):
 
     split_tree = split_tree.at[node_to_prune].set(0)
     affluence_tree = (
-        None if affluence_tree is None else
-        affluence_tree.at[node_to_prune].set(True)
+        None if affluence_tree is None else affluence_tree.at[node_to_prune].set(True)
     )
     is_growable_leaf = growable_leaves(split_tree, affluence_tree)
     prob_choose = p_propose_grow[node_to_prune]
@@ -844,6 +915,7 @@ def choose_leaf_parent(split_tree, affluence_tree, p_propose_grow, key):
 
     return node_to_prune, num_prunable, prob_choose
 
+
 def randint_masked(key, mask):
     """
     Return a random integer in a range, including only some values.
@@ -865,40 +937,46 @@ def randint_masked(key, mask):
     u = random.randint(key, (), 0, ecdf[-1])
     return jnp.searchsorted(ecdf, u, 'right')
 
-def accept_moves_and_sample_leaves(bart, moves, key):
+
+def accept_moves_and_sample_leaves(key, bart, moves):
     """
     Accept or reject the proposed moves and sample the new leaf values.
 
     Parameters
     ----------
+    key : jax.dtypes.prng_key array
+        A jax random key.
     bart : dict
         A BART mcmc state.
     moves : dict
         The proposed moves, see `sample_moves`.
-    key : jax.dtypes.prng_key array
-        A jax random key.
 
     Returns
     -------
     bart : dict
         The new BART mcmc state.
     """
-    bart, moves, count_trees, move_counts, prelkv, prelk, prelf = accept_moves_parallel_stage(bart, moves, key)
-    bart, moves = accept_moves_sequential_stage(bart, count_trees, moves, move_counts, prelkv, prelk, prelf)
+    bart, moves, prec_trees, move_counts, move_precs, prelkv, prelk, prelf = (
+        accept_moves_parallel_stage(key, bart, moves)
+    )
+    bart, moves = accept_moves_sequential_stage(
+        bart, prec_trees, moves, move_counts, move_precs, prelkv, prelk, prelf
+    )
     return accept_moves_final_stage(bart, moves)
 
-def accept_moves_parallel_stage(bart, moves, key):
+
+def accept_moves_parallel_stage(key, bart, moves):
     """
     Pre-computes quantities used to accept moves, in parallel across trees.
 
     Parameters
     ----------
+    key : jax.dtypes.prng_key array
+        A jax random key.
     bart : dict
         A BART mcmc state.
     moves : dict
         The proposed moves, see `sample_moves`.
-    key : jax.dtypes.prng_key array
-        A jax random key.
 
     Returns
     -------
@@ -907,11 +985,14 @@ def accept_moves_parallel_stage(bart, moves, key):
     moves : dict
         The proposed moves, with the field 'partial_ratio' replaced
         by 'log_trans_prior_ratio'.
-    count_trees : array (num_trees, 2 ** d)
-        The number of points in each potential or actual leaf node.
+    prec_trees : float array (num_trees, 2 ** d)
+        The likelihood precision scale in each potential or actual leaf node. If
+        there is no precision scale, this is the number of points in each leaf.
     move_counts : dict
         The counts of the number of points in the the nodes modified by the
         moves.
+    move_precs : dict
+        The likelihood precision scale in each node modified by the moves.
     prelkv, prelk, prelf : dict
         Dictionary with pre-computed terms of the likelihood ratios and leaf
         samples.
@@ -924,20 +1005,35 @@ def accept_moves_parallel_stage(bart, moves, key):
     bart['leaf_trees'] = adapt_leaf_trees_to_grow_indices(bart['leaf_trees'], moves)
 
     # count number of datapoints per leaf
-    count_trees, move_counts = compute_count_trees(bart['leaf_indices'], moves, bart['opt']['count_batch_size'])
+    count_trees, move_counts = compute_count_trees(
+        bart['leaf_indices'], moves, bart['opt']['count_batch_size']
+    )
     if bart['opt']['require_min_points']:
-        count_half_trees = count_trees[:, :bart['var_trees'].shape[1]]
+        count_half_trees = count_trees[:, : bart['var_trees'].shape[1]]
         bart['affluence_trees'] = count_half_trees >= 2 * bart['min_points_per_leaf']
 
+    # count number of datapoints per leaf, weighted by error precision scale
+    if bart['prec_scale'] is None:
+        prec_trees = count_trees
+        move_precs = move_counts
+    else:
+        prec_trees, move_precs = compute_prec_trees(
+            bart['prec_scale'],
+            bart['leaf_indices'],
+            moves,
+            bart['opt']['count_batch_size'],
+        )
+
     # compute some missing information about moves
     moves = complete_ratio(moves, move_counts, bart['min_points_per_leaf'])
     bart['grow_prop_count'] = jnp.sum(moves['grow'])
     bart['prune_prop_count'] = jnp.sum(moves['allowed'] & ~moves['grow'])
 
-    prelkv, prelk = precompute_likelihood_terms(count_trees, bart['sigma2'], move_counts)
-    prelf = precompute_leaf_terms(count_trees, bart['sigma2'], key)
+    prelkv, prelk = precompute_likelihood_terms(bart['sigma2'], move_precs)
+    prelf = precompute_leaf_terms(key, prec_trees, bart['sigma2'])
+
+    return bart, moves, prec_trees, move_counts, move_precs, prelkv, prelk, prelf
 
-    return bart, moves, count_trees, move_counts, prelkv, prelk, prelf
 
 @functools.partial(jaxext.vmap_nodoc, in_axes=(0, 0, None))
 def apply_grow_to_indices(moves, leaf_indices, X):
@@ -968,15 +1064,16 @@ def apply_grow_to_indices(moves, leaf_indices, X):
         leaf_indices,
     )
 
+
 def compute_count_trees(leaf_indices, moves, batch_size):
     """
     Count the number of datapoints in each leaf.
 
     Parameters
     ----------
-    grow_leaf_indices : int array (num_trees, n)
-        The index of the leaf each datapoint falls into, if the grow move is
-        accepted.
+    leaf_indices : int array (num_trees, n)
+        The index of the leaf each datapoint falls into, with the deeper version
+        of the tree (post-GROW, pre-PRUNE).
     moves : dict
         The proposed moves, see `sample_moves`.
     batch_size : int or None
@@ -987,9 +1084,8 @@ def compute_count_trees(leaf_indices, moves, batch_size):
     count_trees : int array (num_trees, 2 ** (d - 1))
         The number of points in each potential or actual leaf node.
     counts : dict
-        The counts of the number of points in the the nodes modified by the
-        moves, organized as two dictionaries 'grow' and 'prune', with subfields
-        'left', 'right', and 'total'.
+        The counts of the number of points in the leaves grown or pruned by the
+        moves, under keys 'left', 'right', and 'total' (left + right).
     """
 
     ntree, tree_size = moves['var_trees'].shape
@@ -1009,6 +1105,7 @@ def compute_count_trees(leaf_indices, moves, batch_size):
 
     return count_trees, counts
 
+
 def count_datapoints_per_leaf(leaf_indices, tree_size, batch_size):
     """
     Count the number of datapoints in each leaf.
@@ -1032,40 +1129,129 @@ def count_datapoints_per_leaf(leaf_indices, tree_size, batch_size):
     else:
         return _count_vec(leaf_indices, tree_size, batch_size)
 
+
 def _count_scan(leaf_indices, tree_size):
     def loop(_, leaf_indices):
         return None, _aggregate_scatter(1, leaf_indices, tree_size, jnp.uint32)
+
     _, count_trees = lax.scan(loop, None, leaf_indices)
     return count_trees
 
+
 def _aggregate_scatter(values, indices, size, dtype):
-    return (jnp
-        .zeros(size, dtype)
-        .at[indices]
-        .add(values)
-    )
+    return jnp.zeros(size, dtype).at[indices].add(values)
+
 
 def _count_vec(leaf_indices, tree_size, batch_size):
-    return _aggregate_batched_alltrees(1, leaf_indices, tree_size, jnp.uint32, batch_size)
-        # uint16 is super-slow on gpu, don't use it even if n < 2^16
+    return _aggregate_batched_alltrees(
+        1, leaf_indices, tree_size, jnp.uint32, batch_size
+    )
+    # uint16 is super-slow on gpu, don't use it even if n < 2^16
+
 
 def _aggregate_batched_alltrees(values, indices, size, dtype, batch_size):
     ntree, n = indices.shape
     tree_indices = jnp.arange(ntree)
     nbatches = n // batch_size + bool(n % batch_size)
     batch_indices = jnp.arange(n) % nbatches
-    return (jnp
-        .zeros((ntree, size, nbatches), dtype)
+    return (
+        jnp.zeros((ntree, size, nbatches), dtype)
         .at[tree_indices[:, None], indices, batch_indices]
         .add(values)
         .sum(axis=2)
     )
 
+
+def compute_prec_trees(prec_scale, leaf_indices, moves, batch_size):
+    """
+    Compute the likelihood precision scale in each leaf.
+
+    Parameters
+    ----------
+    prec_scale : float array (n,)
+        The scale of the precision of the error on each datapoint.
+    leaf_indices : int array (num_trees, n)
+        The index of the leaf each datapoint falls into, with the deeper version
+        of the tree (post-GROW, pre-PRUNE).
+    moves : dict
+        The proposed moves, see `sample_moves`.
+    batch_size : int or None
+        The data batch size to use for the summation.
+
+    Returns
+    -------
+    prec_trees : float array (num_trees, 2 ** (d - 1))
+        The likelihood precision scale in each potential or actual leaf node.
+    counts : dict
+        The likelihood precision scale in the leaves grown or pruned by the
+        moves, under keys 'left', 'right', and 'total' (left + right).
+    """
+
+    ntree, tree_size = moves['var_trees'].shape
+    tree_size *= 2
+    tree_indices = jnp.arange(ntree)
+
+    prec_trees = prec_per_leaf(prec_scale, leaf_indices, tree_size, batch_size)
+
+    # prec datapoints in nodes modified by move
+    precs = dict()
+    precs['left'] = prec_trees[tree_indices, moves['left']]
+    precs['right'] = prec_trees[tree_indices, moves['right']]
+    precs['total'] = precs['left'] + precs['right']
+
+    # write prec into non-leaf node
+    prec_trees = prec_trees.at[tree_indices, moves['node']].set(precs['total'])
+
+    return prec_trees, precs
+
+
+def prec_per_leaf(prec_scale, leaf_indices, tree_size, batch_size):
+    """
+    Compute the likelihood precision scale in each leaf.
+
+    Parameters
+    ----------
+    prec_scale : float array (n,)
+        The scale of the precision of the error on each datapoint.
+    leaf_indices : int array (num_trees, n)
+        The index of the leaf each datapoint falls into.
+    tree_size : int
+        The size of the leaf tree array (2 ** d).
+    batch_size : int or None
+        The data batch size to use for the summation.
+
+    Returns
+    -------
+    prec_trees : int array (num_trees, 2 ** (d - 1))
+        The likelihood precision scale in each leaf node.
+    """
+    if batch_size is None:
+        return _prec_scan(prec_scale, leaf_indices, tree_size)
+    else:
+        return _prec_vec(prec_scale, leaf_indices, tree_size, batch_size)
+
+
+def _prec_scan(prec_scale, leaf_indices, tree_size):
+    def loop(_, leaf_indices):
+        return None, _aggregate_scatter(
+            prec_scale, leaf_indices, tree_size, jnp.float32
+        )  # TODO: use large_float
+
+    _, prec_trees = lax.scan(loop, None, leaf_indices)
+    return prec_trees
+
+
+def _prec_vec(prec_scale, leaf_indices, tree_size, batch_size):
+    return _aggregate_batched_alltrees(
+        prec_scale, leaf_indices, tree_size, jnp.float32, batch_size
+    )  # TODO: use large_float
+
+
 def complete_ratio(moves, move_counts, min_points_per_leaf):
     """
     Complete non-likelihood MH ratio calculation.
 
-    This functions adds the probability of choosing the prune move.
+    This function adds the probability of choosing the prune move.
 
     Parameters
     ----------
@@ -1084,10 +1270,13 @@ def complete_ratio(moves, move_counts, min_points_per_leaf):
         'log_trans_prior_ratio'.
     """
     moves = moves.copy()
-    p_prune = compute_p_prune(moves, move_counts['left'], move_counts['right'], min_points_per_leaf)
+    p_prune = compute_p_prune(
+        moves, move_counts['left'], move_counts['right'], min_points_per_leaf
+    )
     moves['log_trans_prior_ratio'] = jnp.log(moves.pop('partial_ratio') * p_prune)
     return moves
 
+
 def compute_p_prune(moves, left_count, right_count, min_points_per_leaf):
     """
     Compute the probability of proposing a prune move.
@@ -1123,6 +1312,7 @@ def compute_p_prune(moves, left_count, right_count, min_points_per_leaf):
 
     return jnp.where(moves['grow'], grow_p_prune, prune_p_prune)
 
+
 @jaxext.vmap_nodoc
 def adapt_leaf_trees_to_grow_indices(leaf_trees, moves):
     """
@@ -1143,26 +1333,25 @@ def adapt_leaf_trees_to_grow_indices(leaf_trees, moves):
         what would be its children if the grow move was accepted.
     """
     values_at_node = leaf_trees[moves['node']]
-    return (leaf_trees
-        .at[jnp.where(moves['grow'], moves['left'], leaf_trees.size)]
+    return (
+        leaf_trees.at[jnp.where(moves['grow'], moves['left'], leaf_trees.size)]
         .set(values_at_node)
         .at[jnp.where(moves['grow'], moves['right'], leaf_trees.size)]
         .set(values_at_node)
     )
 
-def precompute_likelihood_terms(count_trees, sigma2, move_counts):
+
+def precompute_likelihood_terms(sigma2, move_precs):
     """
     Pre-compute terms used in the likelihood ratio of the acceptance step.
 
     Parameters
     ----------
-    count_trees : array (num_trees, 2 ** d)
-        The number of points in each potential or actual leaf node.
     sigma2 : float
         The noise variance.
-    move_counts : dict
-        The counts of the number of points in the the nodes modified by the
-        moves.
+    move_precs : dict
+        The likelihood precision scale in the leaves grown or pruned by the
+        moves, under keys 'left', 'right', and 'total' (left + right).
 
     Returns
     -------
@@ -1173,32 +1362,37 @@ def precompute_likelihood_terms(count_trees, sigma2, move_counts):
         Dictionary with pre-computed terms of the likelihood ratio, shared by
         all trees.
     """
-    ntree = len(count_trees)
+    ntree = len(move_precs['total'])
     sigma_mu2 = 1 / ntree
     prelkv = dict()
-    prelkv['sigma2_left'] = sigma2 + move_counts['left'] * sigma_mu2
-    prelkv['sigma2_right'] = sigma2 + move_counts['right'] * sigma_mu2
-    prelkv['sigma2_total'] = sigma2 + move_counts['total'] * sigma_mu2
-    prelkv['sqrt_term'] = jnp.log(
-        sigma2 * prelkv['sigma2_total'] /
-        (prelkv['sigma2_left'] * prelkv['sigma2_right'])
-    ) / 2
+    prelkv['sigma2_left'] = sigma2 + move_precs['left'] * sigma_mu2
+    prelkv['sigma2_right'] = sigma2 + move_precs['right'] * sigma_mu2
+    prelkv['sigma2_total'] = sigma2 + move_precs['total'] * sigma_mu2
+    prelkv['sqrt_term'] = (
+        jnp.log(
+            sigma2
+            * prelkv['sigma2_total']
+            / (prelkv['sigma2_left'] * prelkv['sigma2_right'])
+        )
+        / 2
+    )
     return prelkv, dict(
         exp_factor=sigma_mu2 / (2 * sigma2),
     )
 
-def precompute_leaf_terms(count_trees, sigma2, key):
+
+def precompute_leaf_terms(key, prec_trees, sigma2):
     """
     Pre-compute terms used to sample leaves from their posterior.
 
     Parameters
     ----------
-    count_trees : array (num_trees, 2 ** d)
-        The number of points in each potential or actual leaf node.
-    sigma2 : float
-        The noise variance.
     key : jax.dtypes.prng_key array
         A jax random key.
+    prec_trees : array (num_trees, 2 ** d)
+        The likelihood precision scale in each potential or actual leaf node.
+    sigma2 : float
+        The noise variance.
 
     Returns
     -------
@@ -1206,22 +1400,25 @@ def precompute_leaf_terms(count_trees, sigma2, key):
         Dictionary with pre-computed terms of the leaf sampling, with fields:
 
         'mean_factor' : float array (num_trees, 2 ** d)
-            The factor to be multiplied by the sum of residuals to obtain the
-            posterior mean.
+            The factor to be multiplied by the sum of the scaled residuals to
+            obtain the posterior mean.
         'centered_leaves' : float array (num_trees, 2 ** d)
             The mean-zero normal values to be added to the posterior mean to
             obtain the posterior leaf samples.
     """
-    ntree = len(count_trees)
-    prec_lk = count_trees / sigma2
-    var_post = lax.reciprocal(prec_lk + ntree) # = 1 / (prec_lk + prec_prior)
-    z = random.normal(key, count_trees.shape, sigma2.dtype)
+    ntree = len(prec_trees)
+    prec_lk = prec_trees / sigma2
+    var_post = lax.reciprocal(prec_lk + ntree)  # = 1 / (prec_lk + prec_prior)
+    z = random.normal(key, prec_trees.shape, sigma2.dtype)
     return dict(
-        mean_factor=var_post / sigma2, # = mean_lk * prec_lk * var_post / resid_tree
+        mean_factor=var_post / sigma2,  # = mean_lk * prec_lk * var_post / resid_tree
         centered_leaves=z * jnp.sqrt(var_post),
     )
 
-def accept_moves_sequential_stage(bart, count_trees, moves, move_counts, prelkv, prelk, prelf):
+
+def accept_moves_sequential_stage(
+    bart, prec_trees, moves, move_counts, move_precs, prelkv, prelk, prelf
+):
     """
     The part of accepting the moves that has to be done one tree at a time.
 
@@ -1229,13 +1426,15 @@ def accept_moves_sequential_stage(bart, count_trees, moves, move_counts, prelkv,
     ----------
     bart : dict
         A partially updated BART mcmc state.
-    count_trees : array (num_trees, 2 ** d)
-        The number of points in each potential or actual leaf node.
+    prec_trees : float array (num_trees, 2 ** d)
+        The likelihood precision scale in each potential or actual leaf node.
     moves : dict
         The proposed moves, see `sample_moves`.
     move_counts : dict
         The counts of the number of points in the the nodes modified by the
         moves.
+    move_precs : dict
+        The likelihood precision scale in each node modified by the moves.
     prelkv, prelk, prelf : dict
         Dictionaries with pre-computed terms of the likelihood ratios and leaf
         samples.
@@ -1262,6 +1461,7 @@ def accept_moves_sequential_stage(bart, count_trees, moves, move_counts, prelkv,
             len(bart['leaf_trees']),
             bart['opt']['resid_batch_size'],
             resid,
+            bart['prec_scale'],
             bart['min_points_per_leaf'],
             'ratios' in bart,
             prelk,
@@ -1270,22 +1470,44 @@ def accept_moves_sequential_stage(bart, count_trees, moves, move_counts, prelkv,
         return resid, (leaf_tree, acc, to_prune, ratios)
 
     items = (
-        bart['leaf_trees'], count_trees,
-        moves, move_counts,
+        bart['leaf_trees'],
+        prec_trees,
+        moves,
+        move_counts,
+        move_precs,
         bart['leaf_indices'],
-        prelkv, prelf,
+        prelkv,
+        prelf,
     )
     resid, (leaf_trees, acc, to_prune, ratios) = lax.scan(loop, bart['resid'], items)
 
     bart['resid'] = resid
     bart['leaf_trees'] = leaf_trees
-    bart.get('ratios', {}).update(ratios)
+    bart.get('ratios', {}).update(ratios)  # noop if there are no ratios
     moves['acc'] = acc
     moves['to_prune'] = to_prune
 
     return bart, moves
 
-def accept_move_and_sample_leaves(X, ntree, resid_batch_size, resid, min_points_per_leaf, save_ratios, prelk, leaf_tree, count_tree, move, move_counts, leaf_indices, prelkv, prelf):
+
+def accept_move_and_sample_leaves(
+    X,
+    ntree,
+    resid_batch_size,
+    resid,
+    prec_scale,
+    min_points_per_leaf,
+    save_ratios,
+    prelk,
+    leaf_tree,
+    prec_tree,
+    move,
+    move_counts,
+    move_precs,
+    leaf_indices,
+    prelkv,
+    prelf,
+):
     """
     Accept or reject a proposed move and sample the new leaf values.
 
@@ -1299,6 +1521,9 @@ def accept_move_and_sample_leaves(X, ntree, resid_batch_size, resid, min_points_
         The batch size for computing the sum of residuals in each leaf.
     resid : float array (n,)
         The residuals (data minus forest value).
+    prec_scale : float array (n,) or None
+        The scale of the precision of the error on each datapoint. If None, it
+        is assumed to be 1.
     min_points_per_leaf : int or None
         The minimum number of data points in a leaf node.
     save_ratios : bool
@@ -1308,10 +1533,15 @@ def accept_move_and_sample_leaves(X, ntree, resid_batch_size, resid, min_points_
         trees.
     leaf_tree : float array (2 ** d,)
         The leaf values of the tree.
-    count_tree : int array (2 ** d,)
-        The number of datapoints in each leaf.
+    prec_tree : float array (2 ** d,)
+        The likelihood precision scale in each potential or actual leaf node.
     move : dict
         The proposed move, see `sample_moves`.
+    move_counts : dict
+        The counts of the number of points in the the nodes modified by the
+        moves.
+    move_precs : dict
+        The likelihood precision scale in each node modified by the moves.
     leaf_indices : int array (n,)
         The leaf indices for the largest version of the tree compatible with
         the move.
@@ -1334,11 +1564,15 @@ def accept_move_and_sample_leaves(X, ntree, resid_batch_size, resid, min_points_
         The acceptance ratios for the moves. Empty if not to be saved.
     """
 
-    # sum residuals and count units per leaf, in tree proposed by grow move
-    resid_tree = sum_resid(resid, leaf_indices, leaf_tree.size, resid_batch_size)
+    # sum residuals in each leaf, in tree proposed by grow move
+    if prec_scale is None:
+        scaled_resid = resid
+    else:
+        scaled_resid = resid * prec_scale
+    resid_tree = sum_resid(scaled_resid, leaf_indices, leaf_tree.size, resid_batch_size)
 
     # subtract starting tree from function
-    resid_tree += count_tree * leaf_tree
+    resid_tree += prec_tree * leaf_tree
 
     # get indices of move
     node = move['node']
@@ -1353,7 +1587,9 @@ def accept_move_and_sample_leaves(X, ntree, resid_batch_size, resid, min_points_
     resid_tree = resid_tree.at[node].set(resid_total)
 
     # compute acceptance ratio
-    log_lk_ratio = compute_likelihood_ratio(resid_total, resid_left, resid_right, prelkv, prelk)
+    log_lk_ratio = compute_likelihood_ratio(
+        resid_total, resid_left, resid_right, prelkv, prelk
+    )
     log_ratio = move['log_trans_prior_ratio'] + log_lk_ratio
     log_ratio = jnp.where(move['grow'], log_ratio, -log_ratio)
     ratios = {}
@@ -1374,10 +1610,10 @@ def accept_move_and_sample_leaves(X, ntree, resid_batch_size, resid, min_points_
     mean_post = resid_tree * prelf['mean_factor']
     leaf_tree = mean_post + prelf['centered_leaves']
 
-    # copy leaves around such that the leaf indices select the right leaf
+    # copy leaves around such that the leaf indices point to the correct leaf
     to_prune = acc ^ move['grow']
-    leaf_tree = (leaf_tree
-        .at[jnp.where(to_prune, left, leaf_tree.size)]
+    leaf_tree = (
+        leaf_tree.at[jnp.where(to_prune, left, leaf_tree.size)]
         .set(leaf_tree[node])
         .at[jnp.where(to_prune, right, leaf_tree.size)]
         .set(leaf_tree[node])
@@ -1388,14 +1624,16 @@ def accept_move_and_sample_leaves(X, ntree, resid_batch_size, resid, min_points_
 
     return resid, leaf_tree, acc, to_prune, ratios
 
-def sum_resid(resid, leaf_indices, tree_size, batch_size):
+
+def sum_resid(scaled_resid, leaf_indices, tree_size, batch_size):
     """
     Sum the residuals in each leaf.
 
     Parameters
     ----------
-    resid : float array (n,)
-        The residuals (data minus forest value).
+    scaled_resid : float array (n,)
+        The residuals (data minus forest value) multiplied by the error
+        precision scale.
     leaf_indices : int array (n,)
         The leaf indices of the tree (in which leaf each data point falls into).
     tree_size : int
@@ -1413,29 +1651,32 @@ def sum_resid(resid, leaf_indices, tree_size, batch_size):
         aggr_func = _aggregate_scatter
     else:
         aggr_func = functools.partial(_aggregate_batched_onetree, batch_size=batch_size)
-    return aggr_func(resid, leaf_indices, tree_size, jnp.float32)
+    return aggr_func(
+        scaled_resid, leaf_indices, tree_size, jnp.float32
+    )  # TODO: use large_float
+
 
 def _aggregate_batched_onetree(values, indices, size, dtype, batch_size):
-    n, = indices.shape
+    (n,) = indices.shape
     nbatches = n // batch_size + bool(n % batch_size)
     batch_indices = jnp.arange(n) % nbatches
-    return (jnp
-        .zeros((size, nbatches), dtype)
+    return (
+        jnp.zeros((size, nbatches), dtype)
         .at[indices, batch_indices]
         .add(values)
         .sum(axis=1)
     )
 
+
 def compute_likelihood_ratio(total_resid, left_resid, right_resid, prelkv, prelk):
     """
     Compute the likelihood ratio of a grow move.
 
     Parameters
     ----------
-    total_resid : float
-        The sum of the residuals in the leaf to grow.
-    left_resid, right_resid : float
-        The sum of the residuals in the left/right child of the leaf to grow.
+    total_resid, left_resid, right_resid : float
+        The sum of the residuals (scaled by error precision scale) of the
+        datapoints falling in the nodes involved in the moves.
     prelkv, prelk : dict
         The pre-computed terms of the likelihood ratio, see
         `precompute_likelihood_terms`.
@@ -1446,12 +1687,13 @@ def compute_likelihood_ratio(total_resid, left_resid, right_resid, prelkv, prelk
         The likelihood ratio P(data | new tree) / P(data | old tree).
     """
     exp_term = prelk['exp_factor'] * (
-        left_resid * left_resid / prelkv['sigma2_left'] +
-        right_resid * right_resid / prelkv['sigma2_right'] -
-        total_resid * total_resid / prelkv['sigma2_total']
+        left_resid * left_resid / prelkv['sigma2_left']
+        + right_resid * right_resid / prelkv['sigma2_right']
+        - total_resid * total_resid / prelkv['sigma2_total']
     )
     return prelkv['sqrt_term'] + exp_term
 
+
 def accept_moves_final_stage(bart, moves):
     """
     The final part of accepting the moves, in parallel across trees.
@@ -1478,7 +1720,8 @@ def accept_moves_final_stage(bart, moves):
     bart['split_trees'] = apply_moves_to_split_trees(bart['split_trees'], moves)
     return bart
 
-@jax.vmap
+
+@jaxext.vmap_nodoc
 def apply_moves_to_leaf_indices(leaf_indices, moves):
     """
     Update the leaf indices to match the accepted move.
@@ -1497,7 +1740,7 @@ def apply_moves_to_leaf_indices(leaf_indices, moves):
     leaf_indices : int array (num_trees, n)
         The updated leaf indices.
     """
-    mask = ~jnp.array(1, leaf_indices.dtype) # ...1111111110
+    mask = ~jnp.array(1, leaf_indices.dtype)  # ...1111111110
     is_child = (leaf_indices & mask) == moves['left']
     return jnp.where(
         is_child & moves['to_prune'],
@@ -1505,7 +1748,8 @@ def apply_moves_to_leaf_indices(leaf_indices, moves):
         leaf_indices,
     )
 
-@jax.vmap
+
+@jaxext.vmap_nodoc
 def apply_moves_to_split_trees(split_trees, moves):
     """
     Update the split trees to match the accepted move.
@@ -1523,31 +1767,36 @@ def apply_moves_to_split_trees(split_trees, moves):
     split_trees : int array (num_trees, 2 ** (d - 1))
         The updated split trees.
     """
-    return (split_trees
-        .at[jnp.where(
-            moves['grow'],
-            moves['node'],
-            split_trees.size,
-        )]
+    return (
+        split_trees.at[
+            jnp.where(
+                moves['grow'],
+                moves['node'],
+                split_trees.size,
+            )
+        ]
         .set(moves['grow_split'].astype(split_trees.dtype))
-        .at[jnp.where(
-            moves['to_prune'],
-            moves['node'],
-            split_trees.size,
-        )]
+        .at[
+            jnp.where(
+                moves['to_prune'],
+                moves['node'],
+                split_trees.size,
+            )
+        ]
         .set(0)
     )
 
-def sample_sigma(bart, key):
+
+def sample_sigma(key, bart):
     """
     Noise variance sampling step of BART MCMC.
 
     Parameters
     ----------
-    bart : dict
-        A BART mcmc state, as created by `init`.
     key : jax.dtypes.prng_key array
         A jax random key.
+    bart : dict
+        A BART mcmc state, as created by `init`.
 
     Returns
     -------
@@ -1558,7 +1807,11 @@ def sample_sigma(bart, key):
 
     resid = bart['resid']
     alpha = bart['sigma2_alpha'] + resid.size / 2
-    norm2 = jnp.dot(resid, resid, preferred_element_type=bart['opt']['large_float'])
+    if bart['prec_scale'] is None:
+        scaled_resid = resid
+    else:
+        scaled_resid = resid * bart['prec_scale']
+    norm2 = resid @ scaled_resid
     beta = bart['sigma2_beta'] + norm2 / 2
 
     sample = random.gamma(key, alpha)