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

bartz/mcmcstep.py

@@ -10,10 +10,10 @@
 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 # copies of the Software, and to permit persons to whom the Software is
 # furnished to do so, subject to the following conditions:
-# 
+#
 # The above copyright notice and this permission notice shall be included in all
 # copies or substantial portions of the Software.
-# 
+#
 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
@@ -34,6 +34,7 @@ range of possible values.
 """
 
 import functools
+import math
 
 import jax
 from jax import random
@@ -54,7 +55,9 @@ def init(*,
     small_float=jnp.float32,
     large_float=jnp.float32,
     min_points_per_leaf=None,
-    suffstat_batch_size='auto',
+    resid_batch_size='auto',
+    count_batch_size='auto',
+    save_ratios=False,
     ):
     """
     Make a BART posterior sampling MCMC initial state.
@@ -82,10 +85,13 @@ def init(*,
         The dtype for scalars, small arrays, and arrays which require accuracy.
     min_points_per_leaf : int, optional
         The minimum number of data points in a leaf node. 0 if not specified.
-    suffstat_batch_size : int, None, str, default 'auto'
-        The batch size for computing sufficient statistics. `None` for no
-        batching. If 'auto', pick a value based on the device of `y`, or the
-        default device.
+    resid_batch_size, count_batch_sizes : int, None, str, default 'auto'
+        The batch sizes, along datapoints, for summing the residuals and
+        counting the number of datapoints in each leaf. `None` for no batching.
+        If 'auto', pick a value based on the device of `y`, or the default
+        device.
+    save_ratios : bool, default False
+        Whether to save the Metropolis-Hastings ratios.
 
     Returns
     -------
@@ -111,6 +117,8 @@ def init(*,
         'p_nonterminal' : large_float array (d,)
             The probability of a nonterminal node at each depth, padded with a
             zero.
+        'p_propose_grow' : large_float array (2 ** (d - 1),)
+            The unnormalized probability of picking a leaf for a grow proposal.
         'sigma2_alpha' : large_float
             The shape parameter of the inverse gamma prior on the noise variance.
         'sigma2_beta' : large_float
@@ -121,6 +129,8 @@ def init(*,
             The response.
         'X' : int array (p, n)
             The predictors.
+        'leaf_indices' : int array (num_trees, n)
+            The index of the leaf each datapoints falls into, for each tree.
         'min_points_per_leaf' : int or None
             The minimum number of data points in a leaf node.
         'affluence_trees' : bool array (num_trees, 2 ** (d - 1)) or None
@@ -129,8 +139,6 @@ def init(*,
         'opt' : LeafDict
             A dictionary with config values:
 
-            'suffstat_batch_size' : int or None
-                The batch size for computing sufficient statistics.
             'small_float' : dtype
                 The dtype for large arrays used in the algorithm.
             'large_float' : dtype
@@ -138,6 +146,16 @@ def init(*,
                 accuracy.
             'require_min_points' : bool
                 Whether the `min_points_per_leaf` parameter is specified.
+            'resid_batch_size', 'count_batch_size' : int or None
+                The data batch sizes for computing the sufficient statistics.
+        'ratios' : dict, optional
+            If `save_ratios` is True, this field is present. It has the fields:
+
+            'log_trans_prior' : large_float array (num_trees,)
+                The log transition and prior Metropolis-Hastings ratio for the
+                proposed move on each tree.
+            'log_likelihood' : large_float array (num_trees,)
+                The log likelihood ratio.
     """
 
     p_nonterminal = jnp.asarray(p_nonterminal, large_float)
@@ -151,24 +169,28 @@ def init(*,
     small_float = jnp.dtype(small_float)
     large_float = jnp.dtype(large_float)
     y = jnp.asarray(y, small_float)
-    suffstat_batch_size = _choose_suffstat_batch_size(suffstat_batch_size, y)
+    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)
 
     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)),
         split_trees=make_forest(max_depth - 1, max_split.dtype),
         resid=jnp.asarray(y, large_float),
-        sigma2=jnp.ones((), large_float),
+        sigma2=sigma2,
         grow_prop_count=jnp.zeros((), int),
         grow_acc_count=jnp.zeros((), int),
         prune_prop_count=jnp.zeros((), int),
         prune_acc_count=jnp.zeros((), int),
         p_nonterminal=p_nonterminal,
+        p_propose_grow=p_nonterminal[grove.tree_depths(2 ** (max_depth - 1))],
         sigma2_alpha=jnp.asarray(sigma2_alpha, large_float),
         sigma2_beta=jnp.asarray(sigma2_beta, large_float),
         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)),
         min_points_per_leaf=(
             None if min_points_per_leaf is None else
             jnp.asarray(min_points_per_leaf)
@@ -178,37 +200,59 @@ def init(*,
             make_forest(max_depth - 1, bool).at[:, 1].set(y.size >= 2 * min_points_per_leaf)
         ),
         opt=jaxext.LeafDict(
-            suffstat_batch_size=suffstat_batch_size,
             small_float=small_float,
             large_float=large_float,
             require_min_points=min_points_per_leaf is not None,
+            resid_batch_size=resid_batch_size,
+            count_batch_size=count_batch_size,
         ),
     )
 
+    if save_ratios:
+        bart['ratios'] = dict(
+            log_trans_prior=jnp.full(num_trees, jnp.nan),
+            log_likelihood=jnp.full(num_trees, jnp.nan),
+        )
+
     return bart
 
-def _choose_suffstat_batch_size(size, y):
-    if size == 'auto':
+def _choose_suffstat_batch_size(resid_batch_size, count_batch_size, y, forest_size):
+
+    @functools.cache
+    def get_platform():
         try:
             device = y.devices().pop()
         except jax.errors.ConcretizationTypeError:
             device = jax.devices()[0]
         platform = device.platform
+        if platform not in ('cpu', 'gpu'):
+            raise KeyError(f'Unknown platform: {platform}')
+        return platform
 
+    if resid_batch_size == 'auto':
+        platform = get_platform()
+        n = max(1, y.size)
         if platform == 'cpu':
-            return None
-                # maybe I should batch residuals (not counts) for numerical
-                # accuracy, even if it's slower
+            resid_batch_size = 2 ** int(round(math.log2(n / 6))) # n/6
         elif platform == 'gpu':
-            return 128 # 128 is good on A100, and V100 at high n
-                       # 512 is good on T4, and V100 at low n
-        else:
-            raise KeyError(f'Unknown platform: {platform}')
-    
-    elif size is not None:
-        return int(size)
-    
-    return size
+            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':
+        platform = get_platform()
+        if platform == 'cpu':
+            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
+            itemsize = 4
+            min_batch_size = int(math.ceil(forest_size * itemsize * n / max_memory))
+            count_batch_size = max(count_batch_size, min_batch_size)
+            count_batch_size = max(1, count_batch_size)
+
+    return resid_batch_size, count_batch_size
 
 def step(bart, key):
     """
@@ -248,14 +292,11 @@ def sample_trees(bart, key):
 
     Notes
     -----
-    This function zeroes the proposal counters before using them.
+    This function zeroes the proposal counters.
     """
-    bart = bart.copy()
     key, subkey = random.split(key)
-    grow_moves, prune_moves = sample_moves(bart, subkey)
-    bart['var_trees'] = grow_moves['var_tree']
-    grow_leaf_indices = grove.traverse_forest(bart['X'], grow_moves['var_tree'], grow_moves['split_tree'])
-    return accept_moves_and_sample_leaves(bart, grow_moves, prune_moves, grow_leaf_indices, key)
+    moves = sample_moves(bart, subkey)
+    return accept_moves_and_sample_leaves(bart, moves, key)
 
 def sample_moves(bart, key):
     """
@@ -270,21 +311,75 @@ def sample_moves(bart, key):
 
     Returns
     -------
-    grow_moves, prune_moves : dict
-        The proposals for grow and prune moves. See `grow_move` and `prune_move`.
+    moves : dict
+        A dictionary with fields:
+
+        'allowed' : bool array (num_trees,)
+            Whether the move is possible.
+        'grow' : bool array (num_trees,)
+            Whether the move is a grow move or a prune move.
+        'num_growable' : int array (num_trees,)
+            The number of growable leaves in the original tree.
+        'node' : int array (num_trees,)
+            The index of the leaf to grow or node to prune.
+        'left', 'right' : int array (num_trees,)
+            The indices of the children of 'node'.
+        'partial_ratio' : float array (num_trees,)
+            A factor of the Metropolis-Hastings ratio of the move. It lacks
+            the likelihood ratio and the probability of proposing the prune
+            move. If the move is Prune, the ratio is inverted.
+        'grow_var' : int array (num_trees,)
+            The decision axes of the new rules.
+        'grow_split' : int array (num_trees,)
+            The decision boundaries of the new rules.
+        'var_trees' : int array (num_trees, 2 ** (d - 1))
+            The updated decision axes of the trees, valid whatever move.
+        'logu' : float array (num_trees,)
+            The logarithm of a uniform (0, 1] random variable to be used to
+            accept the move. It's in (-oo, 0].
     """
-    key = random.split(key, bart['var_trees'].shape[0])
-    return sample_moves_vmap_trees(bart['var_trees'], bart['split_trees'], bart['affluence_trees'], bart['max_split'], bart['p_nonterminal'], key)
+    ntree = bart['leaf_trees'].shape[0]
+    key = random.split(key, 1 + ntree)
+    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)
+
+    u, logu = random.uniform(key, (2, ntree), bart['opt']['large_float'])
 
-@functools.partial(jaxext.vmap_nodoc, in_axes=(0, 0, 0, None, None, 0))
-def sample_moves_vmap_trees(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, key):
+    # 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)
+
+    # compute children indices
+    node = jnp.where(grow, grow_moves['node'], prune_moves['node'])
+    left = node << 1
+    right = left + 1
+
+    return dict(
+        allowed=grow | prune_moves['allowed'],
+        grow=grow,
+        num_growable=grow_moves['num_growable'],
+        node=node,
+        left=left,
+        right=right,
+        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))
+def _sample_moves_vmap_trees(*args):
+    args, key = args[:-1], args[-1]
     key, key1 = random.split(key)
-    args = var_tree, split_tree, affluence_tree, max_split, p_nonterminal
     grow = grow_move(*args, key)
     prune = prune_move(*args, key1)
     return grow, prune
 
-def grow_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, key):
+def grow_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, p_propose_grow, key):
     """
     Tree structure grow move proposal of BART MCMC.
 
@@ -304,6 +399,8 @@ def grow_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, ke
         The maximum split index for each variable.
     p_nonterminal : array (d,)
         The probability of a nonterminal node at each depth.
+    p_propose_grow : array (2 ** (d - 1),)
+        The unnormalized probability of choosing a leaf to grow.
     key : jax.dtypes.prng_key array
         A jax random key.
 
@@ -312,41 +409,42 @@ def grow_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, ke
     grow_move : dict
         A dictionary with fields:
 
-        'allowed' : bool
-            Whether the move is possible.
+        'num_growable' : int
+            The number of growable leaves.
         'node' : int
-            The index of the leaf to grow.
-        'var_tree' : array (2 ** (d - 1),)
-            The new decision axes of the tree.
-        'split_tree' : array (2 ** (d - 1),)
-            The new decision boundaries of the tree.
+            The index of the leaf to grow. ``2 ** d`` if there are no growable
+            leaves.
+        'var', 'split' : int
+            The decision axis and boundary of the new rule.
         'partial_ratio' : float
             A factor of the Metropolis-Hastings ratio of the move. It lacks
             the likelihood ratio and the probability of proposing the prune
             move.
+        'var_tree' : array (2 ** (d - 1),)
+            The updated decision axes of the tree.
     """
 
     key, key1, key2 = random.split(key, 3)
-    
-    leaf_to_grow, num_growable, num_prunable, allowed = choose_leaf(split_tree, affluence_tree, key)
+
+    leaf_to_grow, num_growable, prob_choose, num_prunable = choose_leaf(split_tree, affluence_tree, p_propose_grow, key)
 
     var = choose_variable(var_tree, split_tree, max_split, leaf_to_grow, key1)
     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_tree = split_tree.at[leaf_to_grow].set(split.astype(split_tree.dtype))
 
-    ratio = compute_partial_ratio(num_growable, num_prunable, p_nonterminal, leaf_to_grow, split_tree)
+    ratio = compute_partial_ratio(prob_choose, num_prunable, p_nonterminal, leaf_to_grow)
 
     return dict(
-        allowed=allowed,
+        num_growable=num_growable,
         node=leaf_to_grow,
+        var=var,
+        split=split,
         partial_ratio=ratio,
         var_tree=var_tree,
-        split_tree=split_tree,
     )
 
-def choose_leaf(split_tree, affluence_tree, key):
+def choose_leaf(split_tree, affluence_tree, p_propose_grow, key):
     """
     Choose a leaf node to grow in a tree.
 
@@ -356,6 +454,8 @@ def choose_leaf(split_tree, affluence_tree, key):
         The splitting points of the tree.
     affluence_tree : bool array (2 ** (d - 1),) or None
         Whether a leaf has enough points to be grown.
+    p_propose_grow : array (2 ** (d - 1),)
+        The unnormalized probability of choosing a leaf to grow.
     key : jax.dtypes.prng_key array
         A jax random key.
 
@@ -366,19 +466,21 @@ def choose_leaf(split_tree, affluence_tree, key):
         ``2 ** d``.
     num_growable : int
         The number of leaf nodes that can be grown.
+    prob_choose : float
+        The normalized probability of choosing the selected leaf.
     num_prunable : int
         The number of leaf parents that could be pruned, after converting the
         selected leaf to a non-terminal node.
-    allowed : bool
-        Whether the grow move is allowed.
     """
-    is_growable, allowed = growable_leaves(split_tree, affluence_tree)
-    leaf_to_grow = randint_masked(key, is_growable)
-    leaf_to_grow = jnp.where(allowed, leaf_to_grow, 2 * split_tree.size)
+    is_growable = growable_leaves(split_tree, affluence_tree)
     num_growable = jnp.count_nonzero(is_growable)
+    distr = jnp.where(is_growable, p_propose_grow, 0)
+    leaf_to_grow, distr_norm = categorical(key, distr)
+    leaf_to_grow = jnp.where(num_growable, leaf_to_grow, 2 * split_tree.size)
+    prob_choose = distr[leaf_to_grow] / distr_norm
     is_parent = grove.is_leaves_parent(split_tree.at[leaf_to_grow].set(1))
     num_prunable = jnp.count_nonzero(is_parent)
-    return leaf_to_grow, num_growable, num_prunable, allowed
+    return leaf_to_grow, num_growable, prob_choose, num_prunable
 
 def growable_leaves(split_tree, affluence_tree):
     """
@@ -397,34 +499,32 @@ def growable_leaves(split_tree, affluence_tree):
         The mask indicating the leaf nodes that can be proposed to grow, i.e.,
         that are not at the bottom level and have at least two times the number
         of minimum points per leaf.
-    allowed : bool
-        Whether the grow move is allowed, i.e., there are growable leaves.
     """
     is_growable = grove.is_actual_leaf(split_tree)
     if affluence_tree is not None:
         is_growable &= affluence_tree
-    return is_growable, jnp.any(is_growable)
+    return is_growable
 
-def randint_masked(key, mask):
+def categorical(key, distr):
     """
-    Return a random integer in a range, including only some values.
+    Return a random integer from an arbitrary distribution.
 
     Parameters
     ----------
     key : jax.dtypes.prng_key array
         A jax random key.
-    mask : bool array (n,)
-        The mask indicating the allowed values.
+    distr : float array (n,)
+        An unnormalized probability distribution.
 
     Returns
     -------
     u : int
-        A random integer in the range ``[0, n)``, and which satisfies
-        ``mask[u] == True``. If all values in the mask are `False`, return `n`.
+        A random integer in the range ``[0, n)``. If all probabilities are zero,
+        return ``n``.
     """
-    ecdf = jnp.cumsum(mask)
-    u = random.randint(key, (), 0, ecdf[-1])
-    return jnp.searchsorted(ecdf, u, 'right')
+    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):
     """
@@ -479,7 +579,7 @@ def fully_used_variables(var_tree, split_tree, max_split, leaf_index):
         filled with `p`. The fill values are not guaranteed to be placed in any
         particular order. Variables may appear more than once.
     """
-    
+
     var_to_ignore = ancestor_variables(var_tree, max_split, leaf_index)
     split_range_vec = jax.vmap(split_range, in_axes=(None, None, None, None, 0))
     l, r = split_range_vec(var_tree, split_tree, max_split, leaf_index, var_to_ignore)
@@ -611,7 +711,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(num_growable, num_prunable, p_nonterminal, leaf_to_grow, new_split_tree):
+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.
 
@@ -626,8 +726,6 @@ def compute_partial_ratio(num_growable, num_prunable, p_nonterminal, leaf_to_gro
         The probability of a nonterminal node at each depth.
     leaf_to_grow : int
         The index of the leaf to grow.
-    new_split_tree : array (2 ** (d - 1),)
-        The splitting points of the tree, after the leaf is grown.
 
     Returns
     -------
@@ -640,6 +738,9 @@ def compute_partial_ratio(num_growable, num_prunable, p_nonterminal, leaf_to_gro
     # the two ratios also contain factors num_available_split *
     # num_available_var, but they cancel out
 
+    # p_prune can't be computed here because it needs the count trees, which are
+    # 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
@@ -647,31 +748,33 @@ def compute_partial_ratio(num_growable, num_prunable, p_nonterminal, leaf_to_gro
 
     p_grow = jnp.where(prune_allowed, 0.5, 1)
 
-    trans_ratio = num_growable / (p_grow * num_prunable)
+    inv_trans_ratio = p_grow * prob_choose * num_prunable
 
-    depth = grove.tree_depths(new_split_tree.size)[leaf_to_grow]
+    depth = grove.tree_depths(2 ** (p_nonterminal.size - 1))[leaf_to_grow]
     p_parent = p_nonterminal[depth]
     cp_children = 1 - p_nonterminal[depth + 1]
     tree_ratio = cp_children * cp_children * p_parent / (1 - p_parent)
 
-    return trans_ratio * tree_ratio
+    return tree_ratio / inv_trans_ratio
 
-def prune_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, key):
+def prune_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, p_propose_grow, key):
     """
     Tree structure prune move proposal of BART MCMC.
 
     Parameters
     ----------
-    var_tree : array (2 ** (d - 1),)
+    var_tree : int array (2 ** (d - 1),)
         The variable indices of the tree.
-    split_tree : array (2 ** (d - 1),)
+    split_tree : int array (2 ** (d - 1),)
         The splitting points of the tree.
     affluence_tree : bool array (2 ** (d - 1),) or None
         Whether a leaf has enough points to be grown.
-    max_split : array (p,)
+    max_split : int array (p,)
         The maximum split index for each variable.
-    p_nonterminal : array (d,)
+    p_nonterminal : float array (d,)
         The probability of a nonterminal node at each depth.
+    p_propose_grow : float array (2 ** (d - 1),)
+        The unnormalized probability of choosing a leaf to grow.
     key : jax.dtypes.prng_key array
         A jax random key.
 
@@ -683,16 +786,16 @@ def prune_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, k
         'allowed' : bool
             Whether the move is possible.
         'node' : int
-            The index of the node to prune.
+            The index of the node to prune. ``2 ** d`` if no node can be pruned.
         'partial_ratio' : float
             A factor of the Metropolis-Hastings ratio of the move. It lacks
             the likelihood ratio and the probability of proposing the prune
             move. This ratio is inverted.
     """
-    node_to_prune, num_prunable, num_growable = choose_leaf_parent(split_tree, affluence_tree, key)
+    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
 
-    ratio = compute_partial_ratio(num_growable, num_prunable, p_nonterminal, node_to_prune, split_tree)
+    ratio = compute_partial_ratio(prob_choose, num_prunable, p_nonterminal, node_to_prune)
 
     return dict(
         allowed=allowed,
@@ -700,7 +803,7 @@ def prune_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, k
         partial_ratio=ratio, # it is inverted in accept_move_and_sample_leaves
     )
 
-def choose_leaf_parent(split_tree, affluence_tree, key):
+def choose_leaf_parent(split_tree, affluence_tree, p_propose_grow, key):
     """
     Pick a non-terminal node with leaf children to prune in a tree.
 
@@ -710,6 +813,8 @@ def choose_leaf_parent(split_tree, affluence_tree, key):
         The splitting points of the tree.
     affluence_tree : bool array (2 ** (d - 1),) or None
         Whether a leaf has enough points to be grown.
+    p_propose_grow : array (2 ** (d - 1),)
+        The unnormalized probability of choosing a leaf to grow.
     key : jax.dtypes.prng_key array
         A jax random key.
 
@@ -717,28 +822,50 @@ def choose_leaf_parent(split_tree, affluence_tree, key):
     -------
     node_to_prune : int
         The index of the node to prune. If ``num_prunable == 0``, return
-        ``split_tree.size``.
+        ``2 ** d``.
     num_prunable : int
         The number of leaf parents that could be pruned.
-    num_growable : int
-        The number of leaf nodes that can be grown, after pruning the chosen
-        node.
+    prob_choose : float
+        The normalized probability of choosing the node to prune for growth.
     """
     is_prunable = grove.is_leaves_parent(split_tree)
-    node_to_prune = randint_masked(key, is_prunable)
     num_prunable = jnp.count_nonzero(is_prunable)
+    node_to_prune = randint_masked(key, is_prunable)
+    node_to_prune = jnp.where(num_prunable, node_to_prune, 2 * split_tree.size)
 
-    pruned_split_tree = split_tree.at[node_to_prune].set(0)
-    pruned_affluence_tree = (
+    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)
     )
-    is_growable_leaf, _ = growable_leaves(pruned_split_tree, pruned_affluence_tree)
-    num_growable = jnp.count_nonzero(is_growable_leaf)
+    is_growable_leaf = growable_leaves(split_tree, affluence_tree)
+    prob_choose = p_propose_grow[node_to_prune]
+    prob_choose /= jnp.sum(p_propose_grow, where=is_growable_leaf)
 
-    return node_to_prune, num_prunable, num_growable
+    return node_to_prune, num_prunable, prob_choose
 
-def accept_moves_and_sample_leaves(bart, grow_moves, prune_moves, grow_leaf_indices, key):
+def randint_masked(key, mask):
+    """
+    Return a random integer in a range, including only some values.
+
+    Parameters
+    ----------
+    key : jax.dtypes.prng_key array
+        A jax random key.
+    mask : bool array (n,)
+        The mask indicating the allowed values.
+
+    Returns
+    -------
+    u : int
+        A random integer in the range ``[0, n)``, and which satisfies
+        ``mask[u] == True``. If all values in the mask are `False`, return `n`.
+    """
+    ecdf = jnp.cumsum(mask)
+    u = random.randint(key, (), 0, ecdf[-1])
+    return jnp.searchsorted(ecdf, u, 'right')
+
+def accept_moves_and_sample_leaves(bart, moves, key):
     """
     Accept or reject the proposed moves and sample the new leaf values.
 
@@ -746,14 +873,8 @@ def accept_moves_and_sample_leaves(bart, grow_moves, prune_moves, grow_leaf_indi
     ----------
     bart : dict
         A BART mcmc state.
-    grow_moves : dict
-        The proposals for grow moves, batched over the first axis. See
-        `grow_move`.
-    prune_moves : dict
-        The proposals for prune moves, batched over the first axis. See
-        `prune_move`.
-    grow_leaf_indices : int array (num_trees, n)
-        The leaf indices of the trees proposed by the grow move.
+    moves : dict
+        The proposed moves, see `sample_moves`.
     key : jax.dtypes.prng_key array
         A jax random key.
 
@@ -762,41 +883,409 @@ def accept_moves_and_sample_leaves(bart, grow_moves, prune_moves, grow_leaf_indi
     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)
+    return accept_moves_final_stage(bart, moves)
+
+def accept_moves_parallel_stage(bart, moves, key):
+    """
+    Pre-computes quantities used to accept moves, in parallel across trees.
+
+    Parameters
+    ----------
+    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
+        A partially updated BART mcmc state.
+    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.
+    move_counts : dict
+        The counts of the number of points in the the nodes modified by the
+        moves.
+    prelkv, prelk, prelf : dict
+        Dictionary with pre-computed terms of the likelihood ratios and leaf
+        samples.
+    """
+    bart = bart.copy()
+
+    # where the move is grow, modify the state like the move was accepted
+    bart['var_trees'] = moves['var_trees']
+    bart['leaf_indices'] = apply_grow_to_indices(moves, bart['leaf_indices'], bart['X'])
+    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'])
+    if bart['opt']['require_min_points']:
+        count_half_trees = count_trees[:, :bart['var_trees'].shape[1]]
+        bart['affluence_trees'] = count_half_trees >= 2 * bart['min_points_per_leaf']
+
+    # 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)
+
+    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):
+    """
+    Update the leaf indices to apply a grow move.
+
+    Parameters
+    ----------
+    moves : dict
+        The proposed moves, see `sample_moves`.
+    leaf_indices : array (num_trees, n)
+        The index of the leaf each datapoint falls into.
+    X : array (p, n)
+        The predictors matrix.
+
+    Returns
+    -------
+    grow_leaf_indices : array (num_trees, n)
+        The updated leaf indices.
+    """
+    left_child = moves['node'].astype(leaf_indices.dtype) << 1
+    go_right = X[moves['grow_var'], :] >= moves['grow_split']
+    tree_size = jnp.array(2 * moves['var_trees'].size)
+    node_to_update = jnp.where(moves['grow'], moves['node'], tree_size)
+    return jnp.where(
+        leaf_indices == node_to_update,
+        left_child + go_right,
+        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.
+    moves : dict
+        The proposed moves, see `sample_moves`.
+    batch_size : int or None
+        The data batch size to use for the summation.
+
+    Returns
+    -------
+    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'.
+    """
+
+    ntree, tree_size = moves['var_trees'].shape
+    tree_size *= 2
+    tree_indices = jnp.arange(ntree)
+
+    count_trees = count_datapoints_per_leaf(leaf_indices, tree_size, batch_size)
+
+    # count datapoints in nodes modified by move
+    counts = dict()
+    counts['left'] = count_trees[tree_indices, moves['left']]
+    counts['right'] = count_trees[tree_indices, moves['right']]
+    counts['total'] = counts['left'] + counts['right']
+
+    # write count into non-leaf node
+    count_trees = count_trees.at[tree_indices, moves['node']].set(counts['total'])
+
+    return count_trees, counts
+
+def count_datapoints_per_leaf(leaf_indices, tree_size, batch_size):
+    """
+    Count the number of datapoints in each leaf.
+
+    Parameters
+    ----------
+    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
+    -------
+    count_trees : int array (num_trees, 2 ** (d - 1))
+        The number of points in each leaf node.
+    """
+    if batch_size is None:
+        return _count_scan(leaf_indices, tree_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)
+    )
+
+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
+
+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)
+        .at[tree_indices[:, None], indices, batch_indices]
+        .add(values)
+        .sum(axis=2)
+    )
+
+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.
+
+    Parameters
+    ----------
+    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.
+    min_points_per_leaf : int or None
+        The minimum number of data points in a leaf node.
+
+    Returns
+    -------
+    moves : dict
+        The updated moves, with the field 'partial_ratio' replaced by
+        'log_trans_prior_ratio'.
+    """
+    moves = moves.copy()
+    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.
+
+    Parameters
+    ----------
+    moves : dict
+        The proposed moves, see `sample_moves`.
+    left_count, right_count : int
+        The number of datapoints in the proposed children of the leaf to grow.
+    min_points_per_leaf : int or None
+        The minimum number of data points in a leaf node.
+
+    Returns
+    -------
+    p_prune : float
+        The probability of proposing a prune move. If grow: after accepting the
+        grow move, if prune: right away.
+    """
+
+    # calculation in case the move is grow
+    other_growable_leaves = moves['num_growable'] >= 2
+    new_leaves_growable = moves['node'] < moves['var_trees'].shape[1] // 2
+    if min_points_per_leaf is not None:
+        any_above_threshold = left_count >= 2 * min_points_per_leaf
+        any_above_threshold |= right_count >= 2 * min_points_per_leaf
+        new_leaves_growable &= any_above_threshold
+    grow_again_allowed = other_growable_leaves | new_leaves_growable
+    grow_p_prune = jnp.where(grow_again_allowed, 0.5, 1)
+
+    # calculation in case the move is prune
+    prune_p_prune = jnp.where(moves['num_growable'], 0.5, 1)
+
+    return jnp.where(moves['grow'], grow_p_prune, prune_p_prune)
+
+@jaxext.vmap_nodoc
+def adapt_leaf_trees_to_grow_indices(leaf_trees, moves):
+    """
+    Modify leaf values such that the indices of the grow moves work on the
+    original tree.
+
+    Parameters
+    ----------
+    leaf_trees : float array (num_trees, 2 ** d)
+        The leaf values.
+    moves : dict
+        The proposed moves, see `sample_moves`.
+
+    Returns
+    -------
+    leaf_trees : float array (num_trees, 2 ** d)
+        The modified leaf values. The value of the leaf to grow is copied to
+        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)]
+        .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):
+    """
+    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.
+
+    Returns
+    -------
+    prelkv : dict
+        Dictionary with pre-computed terms of the likelihood ratio, one per
+        tree.
+    prelk : dict
+        Dictionary with pre-computed terms of the likelihood ratio, shared by
+        all trees.
+    """
+    ntree = len(count_trees)
+    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
+    return prelkv, dict(
+        exp_factor=sigma_mu2 / (2 * sigma2),
+    )
+
+def precompute_leaf_terms(count_trees, sigma2, key):
+    """
+    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.
+
+    Returns
+    -------
+    prelf : dict
+        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.
+        '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)
+    return dict(
+        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):
+    """
+    The part of accepting the moves that has to be done one tree at a time.
+
+    Parameters
+    ----------
+    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.
+    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.
+    prelkv, prelk, prelf : dict
+        Dictionaries with pre-computed terms of the likelihood ratios and leaf
+        samples.
+
+    Returns
+    -------
+    bart : dict
+        A partially updated BART mcmc state.
+    moves : dict
+        The proposed moves, with these additional fields:
+
+        'acc' : bool array (num_trees,)
+            Whether the move was accepted.
+        'to_prune' : bool array (num_trees,)
+            Whether, to reflect the acceptance status of the move, the state
+            should be updated by pruning the leaves involved in the move.
+    """
     bart = bart.copy()
-    def loop(carry, item):
-        resid = carry.pop('resid')
-        resid, carry, trees = accept_move_and_sample_leaves(
+    moves = moves.copy()
+
+    def loop(resid, item):
+        resid, leaf_tree, acc, to_prune, ratios = accept_move_and_sample_leaves(
             bart['X'],
             len(bart['leaf_trees']),
-            bart['opt']['suffstat_batch_size'],
+            bart['opt']['resid_batch_size'],
             resid,
-            bart['sigma2'],
             bart['min_points_per_leaf'],
-            carry,
+            'ratios' in bart,
+            prelk,
             *item,
         )
-        carry['resid'] = resid
-        return carry, trees
-    carry = {
-        k: jnp.zeros_like(bart[k]) for k in
-        ['grow_prop_count', 'prune_prop_count', 'grow_acc_count', 'prune_acc_count']
-    }
-    carry['resid'] = bart['resid']
+        return resid, (leaf_tree, acc, to_prune, ratios)
+
     items = (
-        bart['leaf_trees'],
-        bart['split_trees'],
-        bart['affluence_trees'],
-        grow_moves,
-        prune_moves,
-        grow_leaf_indices,
-        random.split(key, len(bart['leaf_trees'])),
+        bart['leaf_trees'], count_trees,
+        moves, move_counts,
+        bart['leaf_indices'],
+        prelkv, prelf,
     )
-    carry, trees = lax.scan(loop, carry, items)
-    bart.update(carry)
-    bart.update(trees)
-    return bart
+    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)
+    moves['acc'] = acc
+    moves['to_prune'] = to_prune
+
+    return bart, moves
 
-def accept_move_and_sample_leaves(X, ntree, suffstat_batch_size, resid, sigma2, min_points_per_leaf, counts, leaf_tree, split_tree, affluence_tree, grow_move, prune_move, grow_leaf_indices, key):
+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):
     """
     Accept or reject a proposed move and sample the new leaf values.
 
@@ -806,158 +1295,102 @@ def accept_move_and_sample_leaves(X, ntree, suffstat_batch_size, resid, sigma2,
         The predictors.
     ntree : int
         The number of trees in the forest.
-    suffstat_batch_size : int, None
-        The batch size for computing sufficient statistics.
+    resid_batch_size : int, None
+        The batch size for computing the sum of residuals in each leaf.
     resid : float array (n,)
         The residuals (data minus forest value).
-    sigma2 : float
-        The noise variance.
     min_points_per_leaf : int or None
         The minimum number of data points in a leaf node.
-    counts : dict
-        The acceptance counts from the mcmc state dict.
+    save_ratios : bool
+        Whether to save the acceptance ratios.
+    prelk : dict
+        The pre-computed terms of the likelihood ratio which are shared across
+        trees.
     leaf_tree : float array (2 ** d,)
         The leaf values of the tree.
-    split_tree : int array (2 ** (d - 1),)
-        The decision boundaries of the tree.
-    affluence_tree : bool array (2 ** (d - 1),) or None
-        Whether a leaf has enough points to be grown.
-    grow_move : dict
-        The proposal for the grow move. See `grow_move`.
-    prune_move : dict
-        The proposal for the prune move. See `prune_move`.
-    grow_leaf_indices : int array (n,)
-        The leaf indices of the tree proposed by the grow move.
-    key : jax.dtypes.prng_key array
-        A jax random key.
+    count_tree : int array (2 ** d,)
+        The number of datapoints in each leaf.
+    move : dict
+        The proposed move, see `sample_moves`.
+    leaf_indices : int array (n,)
+        The leaf indices for the largest version of the tree compatible with
+        the move.
+    prelkv, prelf : dict
+        The pre-computed terms of the likelihood ratio and leaf sampling which
+        are specific to the tree.
 
     Returns
     -------
     resid : float array (n,)
         The updated residuals (data minus forest value).
-    counts : dict
-        The updated acceptance counts.
-    trees : dict
-        The updated tree arrays.
-    """
-    
-    # compute leaf indices in starting tree
-    grow_node = grow_move['node']
-    grow_left = grow_node << 1
-    grow_right = grow_left + 1
-    leaf_indices = jnp.where(
-        (grow_leaf_indices == grow_left) | (grow_leaf_indices == grow_right),
-        grow_node,
-        grow_leaf_indices,
-    )
+    leaf_tree : float array (2 ** d,)
+        The new leaf values of the tree.
+    acc : bool
+        Whether the move was accepted.
+    to_prune : bool
+        Whether, to reflect the acceptance status of the move, the state should
+        be updated by pruning the leaves involved in the move.
+    ratios : dict
+        The acceptance ratios for the moves. Empty if not to be saved.
+    """
 
-    # compute leaf indices in prune tree
-    prune_node = prune_move['node']
-    prune_left = prune_node << 1
-    prune_right = prune_left + 1
-    prune_leaf_indices = jnp.where(
-        (leaf_indices == prune_left) | (leaf_indices == prune_right),
-        prune_node,
-        leaf_indices,
-    )
+    # 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)
 
     # subtract starting tree from function
-    resid += leaf_tree[leaf_indices]
-
-    # aggregate residuals and count units per leaf
-    grow_resid_tree, grow_count_tree = sufficient_stat(resid, grow_leaf_indices, leaf_tree.size, suffstat_batch_size)
-
-    # compute aggregations in starting tree
-    # I do not zero the children because garbage there does not matter
-    resid_tree = (grow_resid_tree.at[grow_node]
-        .set(grow_resid_tree[grow_left] + grow_resid_tree[grow_right]))
-    count_tree = (grow_count_tree.at[grow_node]
-        .set(grow_count_tree[grow_left] + grow_count_tree[grow_right]))
-
-    # compute aggregations in prune tree
-    prune_resid_tree = (resid_tree.at[prune_node]
-        .set(resid_tree[prune_left] + resid_tree[prune_right]))
-    prune_count_tree = (count_tree.at[prune_node]
-        .set(count_tree[prune_left] + count_tree[prune_right]))
-
-    # compute affluence trees
-    if min_points_per_leaf is not None:
-        grow_affluence_tree = grow_count_tree[:grow_count_tree.size // 2] >= 2 * min_points_per_leaf
-        prune_affluence_tree = affluence_tree.at[prune_node].set(True)
-
-    # compute probability of proposing prune
-    grow_p_prune = compute_p_prune_back(grow_move['split_tree'], grow_affluence_tree)
-    prune_p_prune = compute_p_prune_back(split_tree, affluence_tree)
-
-    # compute likelihood ratios
-    grow_lk_ratio = compute_likelihood_ratio(grow_resid_tree, grow_count_tree, sigma2, grow_node, ntree, min_points_per_leaf)
-    prune_lk_ratio = compute_likelihood_ratio(resid_tree, count_tree, sigma2, prune_node, ntree, min_points_per_leaf)
-
-    # compute acceptance ratios
-    grow_ratio = grow_p_prune * grow_move['partial_ratio'] * grow_lk_ratio
-    prune_ratio = prune_p_prune * prune_move['partial_ratio'] * prune_lk_ratio
-    prune_ratio = lax.reciprocal(prune_ratio)
-
-    # random coins in [0, 1) for proposal and acceptance
-    key, subkey = random.split(key)
-    u0, u1 = random.uniform(subkey, (2,))
-
-    # determine what move to propose (not proposing anything is an option)
-    p_grow = jnp.where(grow_move['allowed'] & prune_move['allowed'], 0.5, grow_move['allowed'])
-    try_grow = u0 < p_grow
-    try_prune = prune_move['allowed'] & ~try_grow
+    resid_tree += count_tree * leaf_tree
+
+    # get indices of move
+    node = move['node']
+    assert node.dtype == jnp.int32
+    left = move['left']
+    right = move['right']
+
+    # sum residuals in parent node modified by move
+    resid_left = resid_tree[left]
+    resid_right = resid_tree[right]
+    resid_total = resid_left + resid_right
+    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_ratio = move['log_trans_prior_ratio'] + log_lk_ratio
+    log_ratio = jnp.where(move['grow'], log_ratio, -log_ratio)
+    ratios = {}
+    if save_ratios:
+        ratios.update(
+            log_trans_prior=move['log_trans_prior_ratio'],
+            log_likelihood=log_lk_ratio,
+        )
 
     # determine whether to accept the move
-    do_grow = try_grow & (u1 < grow_ratio)
-    do_prune = try_prune & (u1 < prune_ratio)
-
-    # pick trees for chosen move
-    trees = {}
-    split_tree = jnp.where(do_grow, grow_move['split_tree'], split_tree)
-    # the prune var tree is equal to the initial one, because I leave garbage values behind
-    split_tree = split_tree.at[prune_node].set(
-        jnp.where(do_prune, 0, split_tree[prune_node]))
+    acc = move['allowed'] & (move['logu'] <= log_ratio)
     if min_points_per_leaf is not None:
-        affluence_tree = jnp.where(do_grow, grow_affluence_tree, affluence_tree)
-        affluence_tree = jnp.where(do_prune, prune_affluence_tree, affluence_tree)
-    resid_tree = jnp.where(do_grow, grow_resid_tree, resid_tree)
-    count_tree = jnp.where(do_grow, grow_count_tree, count_tree)
-    resid_tree = jnp.where(do_prune, prune_resid_tree, resid_tree)
-    count_tree = jnp.where(do_prune, prune_count_tree, count_tree)
-
-    # update acceptance counts
-    counts = counts.copy()
-    counts['grow_prop_count'] += try_grow
-    counts['grow_acc_count'] += do_grow
-    counts['prune_prop_count'] += try_prune
-    counts['prune_acc_count'] += do_prune
-
-    # compute leaves posterior
-    prec_lk = count_tree / sigma2
-    var_post = lax.reciprocal(prec_lk + ntree) # = 1 / (prec_lk + prec_prior)
-    mean_post = resid_tree / sigma2 * var_post # = mean_lk * prec_lk * var_post
-
-    # sample leaves
-    z = random.normal(key, mean_post.shape, mean_post.dtype)
-    leaf_tree = mean_post + z * jnp.sqrt(var_post)
-
-    # add new tree to function
-    leaf_indices = jnp.where(do_grow, grow_leaf_indices, leaf_indices)
-    leaf_indices = jnp.where(do_prune, prune_leaf_indices, leaf_indices)
-    resid -= leaf_tree[leaf_indices]
+        acc &= move_counts['left'] >= min_points_per_leaf
+        acc &= move_counts['right'] >= min_points_per_leaf
+
+    # compute leaves posterior and sample leaves
+    initial_leaf_tree = leaf_tree
+    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
+    to_prune = acc ^ move['grow']
+    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])
+    )
 
-    # pack trees
-    trees = {
-        'leaf_trees': leaf_tree,
-        'split_trees': split_tree,
-        'affluence_trees': affluence_tree,
-    }
+    # replace old tree with new tree in function values
+    resid += (initial_leaf_tree - leaf_tree)[leaf_indices]
 
-    return resid, counts, trees
+    return resid, leaf_tree, acc, to_prune, ratios
 
-def sufficient_stat(resid, leaf_indices, tree_size, batch_size):
+def sum_resid(resid, leaf_indices, tree_size, batch_size):
     """
-    Compute the sufficient statistics for the likelihood ratio of a tree move.
+    Sum the residuals in each leaf.
 
     Parameters
     ----------
@@ -968,124 +1401,142 @@ def sufficient_stat(resid, leaf_indices, tree_size, batch_size):
     tree_size : int
         The size of the tree array (2 ** d).
     batch_size : int, None
-        The batch size for the aggregation. Batching increases numerical
+        The data batch size for the aggregation. Batching increases numerical
         accuracy and parallelism.
 
     Returns
     -------
     resid_tree : float array (2 ** d,)
         The sum of the residuals at data points in each leaf.
-    count_tree : int array (2 ** d,)
-        The number of data points in each leaf.
     """
     if batch_size is None:
         aggr_func = _aggregate_scatter
     else:
-        aggr_func = functools.partial(_aggregate_batched, batch_size=batch_size)
-    resid_tree = aggr_func(resid, leaf_indices, tree_size, jnp.float32)
-    count_tree = aggr_func(1, leaf_indices, tree_size, jnp.uint32)
-    return resid_tree, count_tree
+        aggr_func = functools.partial(_aggregate_batched_onetree, batch_size=batch_size)
+    return aggr_func(resid, leaf_indices, tree_size, jnp.float32)
 
-def _aggregate_scatter(values, indices, size, dtype):
+def _aggregate_batched_onetree(values, indices, size, dtype, batch_size):
+    n, = indices.shape
+    nbatches = n // batch_size + bool(n % batch_size)
+    batch_indices = jnp.arange(n) % nbatches
     return (jnp
-        .zeros(size, dtype)
-        .at[indices]
+        .zeros((size, nbatches), dtype)
+        .at[indices, batch_indices]
         .add(values)
+        .sum(axis=1)
     )
 
-def _aggregate_batched(values, indices, size, dtype, batch_size):
-    nbatches = indices.size // batch_size + bool(indices.size % batch_size)
-    batch_indices = jnp.arange(indices.size) // batch_size
-    return (jnp
-        .zeros((nbatches, size), dtype)
-        .at[batch_indices, indices]
-        .add(values)
-        .sum(axis=0)
-    )
-
-def compute_p_prune_back(new_split_tree, new_affluence_tree):
+def compute_likelihood_ratio(total_resid, left_resid, right_resid, prelkv, prelk):
     """
-    Compute the probability of proposing a prune move after doing a grow move.
+    Compute the likelihood ratio of a grow move.
 
     Parameters
     ----------
-    new_split_tree : int array (2 ** (d - 1),)
-        The decision boundaries of the tree, after the grow move.
-    new_affluence_tree : bool array (2 ** (d - 1),)
-        Which leaves have enough points to be grown, after the grow move.
+    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.
+    prelkv, prelk : dict
+        The pre-computed terms of the likelihood ratio, see
+        `precompute_likelihood_terms`.
 
     Returns
     -------
-    p_prune : float
-        The probability of proposing a prune move after the grow move. This is
-        0.5 if grow is possible again, and 1 if it isn't. It can't be 0 because
-        at least the node just grown can be pruned.
+    ratio : float
+        The likelihood ratio P(data | new tree) / P(data | old tree).
     """
-    _, grow_again_allowed = growable_leaves(new_split_tree, new_affluence_tree)
-    return jnp.where(grow_again_allowed, 0.5, 1)
+    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']
+    )
+    return prelkv['sqrt_term'] + exp_term
 
-def compute_likelihood_ratio(resid_tree, count_tree, sigma2, node, n_tree, min_points_per_leaf):
+def accept_moves_final_stage(bart, moves):
     """
-    Compute the likelihood ratio of a grow move.
+    The final part of accepting the moves, in parallel across trees.
 
     Parameters
     ----------
-    resid_tree : float array (2 ** d,)
-        The sum of the residuals at data points in each leaf.
-    count_tree : int array (2 ** d,)
-        The number of data points in each leaf.
-    sigma2 : float
-        The noise variance.
-    node : int
-        The index of the leaf that has been grown.
-    n_tree : int
-        The number of trees in the forest.
-    min_points_per_leaf : int or None
-        The minimum number of data points in a leaf node.
+    bart : dict
+        A partially updated BART mcmc state.
+    counts : dict
+        The indicators of proposals and acceptances for grow and prune moves.
+    moves : dict
+        The proposed moves (see `sample_moves`) as updated by
+        `accept_moves_sequential_stage`.
 
     Returns
     -------
-    ratio : float
-        The likelihood ratio P(data | new tree) / P(data | old tree).
-
-    Notes
-    -----
-    The ratio is set to 0 if the grow move would create leaves with not enough
-    datapoints per leaf, although this is part of the prior rather than the
-    likelihood.
+    bart : dict
+        The fully updated BART mcmc state.
     """
+    bart = bart.copy()
+    bart['grow_acc_count'] = jnp.sum(moves['acc'] & moves['grow'])
+    bart['prune_acc_count'] = jnp.sum(moves['acc'] & ~moves['grow'])
+    bart['leaf_indices'] = apply_moves_to_leaf_indices(bart['leaf_indices'], moves)
+    bart['split_trees'] = apply_moves_to_split_trees(bart['split_trees'], moves)
+    return bart
 
-    left_child = node << 1
-    right_child = left_child + 1
-
-    left_resid = resid_tree[left_child]
-    right_resid = resid_tree[right_child]
-    total_resid = left_resid + right_resid
-
-    left_count = count_tree[left_child]
-    right_count = count_tree[right_child]
-    total_count = left_count + right_count
-
-    sigma_mu2 = 1 / n_tree
-    sigma2_left = sigma2 + left_count * sigma_mu2
-    sigma2_right = sigma2 + right_count * sigma_mu2
-    sigma2_total = sigma2 + total_count * sigma_mu2
+@jax.vmap
+def apply_moves_to_leaf_indices(leaf_indices, moves):
+    """
+    Update the leaf indices to match the accepted move.
 
-    sqrt_term = sigma2 * sigma2_total / (sigma2_left * sigma2_right)
+    Parameters
+    ----------
+    leaf_indices : int array (num_trees, n)
+        The index of the leaf each datapoint falls into, if the grow move was
+        accepted.
+    moves : dict
+        The proposed moves (see `sample_moves`), as updated by
+        `accept_moves_sequential_stage`.
 
-    exp_term = sigma_mu2 / (2 * sigma2) * (
-        left_resid * left_resid / sigma2_left +
-        right_resid * right_resid / sigma2_right -
-        total_resid * total_resid / sigma2_total
+    Returns
+    -------
+    leaf_indices : int array (num_trees, n)
+        The updated leaf indices.
+    """
+    mask = ~jnp.array(1, leaf_indices.dtype) # ...1111111110
+    is_child = (leaf_indices & mask) == moves['left']
+    return jnp.where(
+        is_child & moves['to_prune'],
+        moves['node'].astype(leaf_indices.dtype),
+        leaf_indices,
     )
 
-    ratio = jnp.sqrt(sqrt_term) * jnp.exp(exp_term)
+@jax.vmap
+def apply_moves_to_split_trees(split_trees, moves):
+    """
+    Update the split trees to match the accepted move.
 
-    if min_points_per_leaf is not None:
-        ratio = jnp.where(right_count >= min_points_per_leaf, ratio, 0)
-        ratio = jnp.where(left_count >= min_points_per_leaf, ratio, 0)
+    Parameters
+    ----------
+    split_trees : int array (num_trees, 2 ** (d - 1))
+        The cutpoints of the decision nodes in the initial trees.
+    moves : dict
+        The proposed moves (see `sample_moves`), as updated by
+        `accept_moves_sequential_stage`.
 
-    return ratio
+    Returns
+    -------
+    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,
+        )]
+        .set(moves['grow_split'].astype(split_trees.dtype))
+        .at[jnp.where(
+            moves['to_prune'],
+            moves['node'],
+            split_trees.size,
+        )]
+        .set(0)
+    )
 
 def sample_sigma(bart, key):
     """
@@ -1107,7 +1558,7 @@ def sample_sigma(bart, key):
 
     resid = bart['resid']
     alpha = bart['sigma2_alpha'] + resid.size / 2
-    norm2 = jnp.dot(resid, resid, preferred_element_type=bart['sigma2_beta'].dtype)
+    norm2 = jnp.dot(resid, resid, preferred_element_type=bart['opt']['large_float'])
     beta = bart['sigma2_beta'] + norm2 / 2
 
     sample = random.gamma(key, alpha)