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

bartz/{interface.py → BART.py}

@@ -1,4 +1,4 @@
-# bartz/src/bartz/interface.py
+# bartz/src/bartz/BART.py
 #
 # Copyright (c) 2024, Giacomo Petrillo
 #
@@ -33,7 +33,7 @@ from . import mcmcstep
 from . import mcmcloop
 from . import prepcovars
 
-class BART:
+class gbart:
     """
     Nonparametric regression with Bayesian Additive Regression Trees (BART).
 
@@ -133,7 +133,7 @@ class BART:
 
     Notes
     -----
-    This interface imitates the function `wbart` from the R package `BART
+    This interface imitates the function `gbart` from the R package `BART
     <https://cran.r-project.org/package=BART>`_, but with these differences:
 
     - If `x_train` and `x_test` are matrices, they have one predictor per row
@@ -142,6 +142,7 @@ class BART:
     - `usequants` is always `True`.
     - `rm_const` is always `False`.
     - The default `numcut` is 255 instead of 100.
+    - A lot of functionality is missing (variable selection, discrete response).
     - There are some additional attributes, and some missing.
     """
 

bartz/__init__.py

@@ -30,6 +30,11 @@ See the manual at https://gattocrucco.github.io/bartz/docs
 
 from ._version import __version__
 
-from .interface import BART
+from . import BART
 
 from . import debug
+from . import grove
+from . import mcmcstep
+from . import mcmcloop
+from . import prepcovars
+from . import jaxext

bartz/_version.py

@@ -1 +1 @@
-__version__ = '0.1.0'
+__version__ = '0.2.1'

bartz/debug.py

@@ -6,6 +6,7 @@ from jax import lax
 
 from . import grove
 from . import mcmcstep
+from . import jaxext
 
 def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
 
@@ -83,7 +84,7 @@ def trace_depth_distr(split_trees_trace):
 def points_per_leaf_distr(var_tree, split_tree, X):
     traverse_tree = jax.vmap(grove.traverse_tree, in_axes=(1, None, None))
     indices = traverse_tree(X, var_tree, split_tree)
-    count_tree = jnp.zeros(2 * split_tree.size, dtype=grove.minimal_unsigned_dtype(indices.size))
+    count_tree = jnp.zeros(2 * split_tree.size, dtype=jaxext.minimal_unsigned_dtype(indices.size))
     count_tree = count_tree.at[indices].add(1)
     is_leaf = grove.is_actual_leaf(split_tree, add_bottom_level=True).view(jnp.uint8)
     return jnp.bincount(count_tree, is_leaf, length=X.shape[1] + 1)
@@ -103,7 +104,7 @@ def trace_points_per_leaf_distr(bart, X):
     return distr
 
 def check_types(leaf_tree, var_tree, split_tree, max_split):
-    expected_var_dtype = grove.minimal_unsigned_dtype(max_split.size - 1)
+    expected_var_dtype = jaxext.minimal_unsigned_dtype(max_split.size - 1)
     expected_split_dtype = max_split.dtype
     return var_tree.dtype == expected_var_dtype and split_tree.dtype == expected_split_dtype
 
@@ -117,7 +118,7 @@ def check_leaf_values(leaf_tree, var_tree, split_tree, max_split):
     return jnp.all(jnp.isfinite(leaf_tree))
 
 def check_stray_nodes(leaf_tree, var_tree, split_tree, max_split):
-    index = jnp.arange(2 * split_tree.size, dtype=grove.minimal_unsigned_dtype(2 * split_tree.size - 1))
+    index = jnp.arange(2 * split_tree.size, dtype=jaxext.minimal_unsigned_dtype(2 * split_tree.size - 1))
     parent_index = index >> 1
     is_not_leaf = split_tree.at[index].get(mode='fill', fill_value=0) != 0
     parent_is_leaf = split_tree[parent_index] == 0
@@ -134,7 +135,7 @@ check_functions = [
 ]
 
 def check_tree(leaf_tree, var_tree, split_tree, max_split):
-    error_type = grove.minimal_unsigned_dtype(2 ** len(check_functions) - 1)
+    error_type = jaxext.minimal_unsigned_dtype(2 ** len(check_functions) - 1)
     error = error_type(0)
     for i, func in enumerate(check_functions):
         ok = func(leaf_tree, var_tree, split_tree, max_split)

bartz/grove.py

@@ -44,7 +44,6 @@ import functools
 import math
 
 import jax
-
 from jax import numpy as jnp
 from jax import lax
 
@@ -107,29 +106,47 @@ def traverse_tree(x, var_tree, split_tree):
 
     carry = (
         jnp.zeros((), bool),
-        jnp.ones((), minimal_unsigned_dtype(2 * var_tree.size - 1)),
+        jnp.ones((), jaxext.minimal_unsigned_dtype(2 * var_tree.size - 1)),
     )
 
     def loop(carry, _):
         leaf_found, index = carry
 
-        split = split_tree.at[index].get(mode='fill', fill_value=0)
-        var = var_tree.at[index].get(mode='fill', fill_value=0)
+        split = split_tree[index]
+        var = var_tree[index]
         
-        leaf_found |= split_tree.at[index].get(mode='fill', fill_value=0) == 0
+        leaf_found |= split == 0
         child_index = (index << 1) + (x[var] >= split)
         index = jnp.where(leaf_found, index, child_index)
 
         return (leaf_found, index), None
 
-        # TODO
-        # - unroll (how much? 5?)
-        # - separate and special-case the last iteration
-
-    depth = 1 + tree_depth(var_tree)
-    (_, index), _ = lax.scan(loop, carry, None, depth)
+    depth = tree_depth(var_tree)
+    (_, index), _ = lax.scan(loop, carry, None, depth, unroll=16)
     return index
 
+@functools.partial(jaxext.vmap_nodoc, in_axes=(None, 0, 0))
+@functools.partial(jaxext.vmap_nodoc, in_axes=(1, None, None))
+def traverse_forest(X, var_trees, split_trees):
+    """
+    Find the leaves where points fall into.
+
+    Parameters
+    ----------
+    X : array (p, n)
+        The coordinates to evaluate the trees at.
+    var_trees : array (m, 2 ** (d - 1))
+        The decision axes of the trees.
+    split_trees : array (m, 2 ** (d - 1))
+        The decision boundaries of the trees.
+
+    Returns
+    -------
+    indices : array (m, n)
+        The indices of the leaves.
+    """
+    return traverse_tree(X, var_trees, split_trees)
+
 def evaluate_forest(X, leaf_trees, var_trees, split_trees, dtype):
     """
     Evaluate a ensemble of trees at an array of points.
@@ -138,7 +155,7 @@ def evaluate_forest(X, leaf_trees, var_trees, split_trees, dtype):
     ----------
     X : array (p, n)
         The coordinates to evaluate the trees at.
-    leaf_trees : (m, 2 ** d)
+    leaf_trees : array (m, 2 ** d)
         The leaf values of the tree or forest. If the input is a forest, the
         first axis is the tree index, and the values are summed.
     var_trees : array (m, 2 ** (d - 1))
@@ -153,30 +170,13 @@ def evaluate_forest(X, leaf_trees, var_trees, split_trees, dtype):
     out : array (n,)
         The sum of the values of the trees at the points in `X`.
     """
-    indices = _traverse_forest(X, var_trees, split_trees)
+    indices = traverse_forest(X, var_trees, split_trees)
     ntree, _ = leaf_trees.shape
-    tree_index = jnp.arange(ntree, dtype=minimal_unsigned_dtype(ntree - 1))[:, None]
+    tree_index = jnp.arange(ntree, dtype=jaxext.minimal_unsigned_dtype(ntree - 1))[:, None]
     leaves = leaf_trees[tree_index, indices]
     return jnp.sum(leaves, axis=0, dtype=dtype)
-        # this sum suggests to swap the vmaps, but I think it's better for X copying to keep it that way
-
-@functools.partial(jax.vmap, in_axes=(None, 0, 0))
-@functools.partial(jax.vmap, in_axes=(1, None, None))
-def _traverse_forest(X, var_trees, split_trees):
-    return traverse_tree(X, var_trees, split_trees)
-
-def minimal_unsigned_dtype(max_value):
-    """
-    Return the smallest unsigned integer dtype that can represent a given
-    maximum value.
-    """
-    if max_value < 2 ** 8:
-        return jnp.uint8
-    if max_value < 2 ** 16:
-        return jnp.uint16
-    if max_value < 2 ** 32:
-        return jnp.uint32
-    return jnp.uint64
+        # this sum suggests to swap the vmaps, but I think it's better for X
+        # copying to keep it that way
 
 def is_actual_leaf(split_tree, *, add_bottom_level=False):
     """
@@ -200,7 +200,7 @@ def is_actual_leaf(split_tree, *, add_bottom_level=False):
     if add_bottom_level:
         size *= 2
         is_leaf = jnp.concatenate([is_leaf, jnp.ones_like(is_leaf)])
-    index = jnp.arange(size, dtype=minimal_unsigned_dtype(size - 1))
+    index = jnp.arange(size, dtype=jaxext.minimal_unsigned_dtype(size - 1))
     parent_index = index >> 1
     parent_nonleaf = split_tree[parent_index].astype(bool)
     parent_nonleaf = parent_nonleaf.at[1].set(True)
@@ -220,7 +220,7 @@ def is_leaves_parent(split_tree):
     is_leaves_parent : bool array (2 ** (d - 1),)
         The mask indicating which nodes have leaf children.
     """
-    index = jnp.arange(split_tree.size, dtype=minimal_unsigned_dtype(2 * split_tree.size - 1))
+    index = jnp.arange(split_tree.size, dtype=jaxext.minimal_unsigned_dtype(2 * split_tree.size - 1))
     left_index = index << 1 # left child
     right_index = left_index + 1 # right child
     left_leaf = split_tree.at[left_index].get(mode='fill', fill_value=0) == 0
@@ -252,4 +252,4 @@ def tree_depths(tree_length):
             depth += 1
         depths.append(depth - 1)
     depths[0] = 0
-    return jnp.array(depths, minimal_unsigned_dtype(max(depths)))
+    return jnp.array(depths, jaxext.minimal_unsigned_dtype(max(depths)))

bartz/jaxext.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
@@ -23,12 +23,19 @@
 # SOFTWARE.
 
 import functools
+import math
+import warnings
 
 from scipy import special
 import jax
 from jax import numpy as jnp
+from jax import tree_util
+from jax import lax
 
 def float_type(*args):
+    """
+    Determine the jax floating point result type given operands/types.
+    """
     t = jnp.result_type(*args)
     return jnp.sin(jnp.empty(0, t)).dtype
 
@@ -39,8 +46,8 @@ def castto(func, type):
     return newfunc
 
 def pure_callback_ufunc(callback, dtype, *args, excluded=None, **kwargs):
-    """ version of jax.pure_callback that deals correctly with ufuncs,
-    see https://github.com/google/jax/issues/17187 """
+    """ version of `jax.pure_callback` that deals correctly with ufuncs,
+    see `<https://github.com/google/jax/issues/17187>`_ """
     if excluded is None:
         excluded = ()
     shape = jnp.broadcast_shapes(*(
@@ -63,6 +70,7 @@ class scipy:
 
     class special:
 
+        @functools.wraps(special.gammainccinv)
         def gammainccinv(a, y):
             a = jnp.asarray(a)
             y = jnp.asarray(y)
@@ -73,13 +81,261 @@ class scipy:
     class stats:
 
         class invgamma:
-            
+
             def ppf(q, a):
                 return 1 / scipy.special.gammainccinv(a, q)
 
 @functools.wraps(jax.vmap)
 def vmap_nodoc(fun, *args, **kw):
+    """
+    Version of `jax.vmap` that preserves the docstring of the input function.
+    """
     doc = fun.__doc__
     fun = jax.vmap(fun, *args, **kw)
     fun.__doc__ = doc
     return fun
+
+def huge_value(x):
+    """
+    Return the maximum value that can be stored in `x`.
+
+    Parameters
+    ----------
+    x : array
+        A numerical numpy or jax array.
+
+    Returns
+    -------
+    maxval : scalar
+        The maximum value allowed by `x`'s type (+inf for floats).
+    """
+    if jnp.issubdtype(x.dtype, jnp.integer):
+        return jnp.iinfo(x.dtype).max
+    else:
+        return jnp.inf
+
+def minimal_unsigned_dtype(max_value):
+    """
+    Return the smallest unsigned integer dtype that can represent a given
+    maximum value (inclusive).
+    """
+    if max_value < 2 ** 8:
+        return jnp.uint8
+    if max_value < 2 ** 16:
+        return jnp.uint16
+    if max_value < 2 ** 32:
+        return jnp.uint32
+    return jnp.uint64
+
+def signed_to_unsigned(int_dtype):
+    """
+    Map a signed integer type to its unsigned counterpart. Unsigned types are
+    passed through.
+    """
+    assert jnp.issubdtype(int_dtype, jnp.integer)
+    if jnp.issubdtype(int_dtype, jnp.unsignedinteger):
+        return int_dtype
+    if int_dtype == jnp.int8:
+        return jnp.uint8
+    if int_dtype == jnp.int16:
+        return jnp.uint16
+    if int_dtype == jnp.int32:
+        return jnp.uint32
+    if int_dtype == jnp.int64:
+        return jnp.uint64
+
+def ensure_unsigned(x):
+    """
+    If x has signed integer type, cast it to the unsigned dtype of the same size.
+    """
+    return x.astype(signed_to_unsigned(x.dtype))
+
+@functools.partial(jax.jit, static_argnums=(1,))
+def unique(x, size, fill_value):
+    """
+    Restricted version of `jax.numpy.unique` that uses less memory.
+
+    Parameters
+    ----------
+    x : 1d array
+        The input array.
+    size : int
+        The length of the output.
+    fill_value : scalar
+        The value to fill the output with if `size` is greater than the number
+        of unique values in `x`.
+
+    Returns
+    -------
+    out : array (size,)
+        The unique values in `x`, sorted, and right-padded with `fill_value`.
+    actual_length : int
+        The number of used values in `out`.
+    """
+    if x.size == 0:
+        return jnp.full(size, fill_value, x.dtype), 0
+    if size == 0:
+        return jnp.empty(0, x.dtype), 0
+    x = jnp.sort(x)
+    def loop(carry, x):
+        i_out, i_in, last, out = carry
+        i_out = jnp.where(x == last, i_out, i_out + 1)
+        out = out.at[i_out].set(x)
+        return (i_out, i_in + 1, x, out), None
+    carry = 0, 0, x[0], jnp.full(size, fill_value, x.dtype)
+    (actual_length, _, _, out), _ = jax.lax.scan(loop, carry, x[:size])
+    return out, actual_length + 1
+
+def autobatch(func, max_io_nbytes, in_axes=0, out_axes=0, return_nbatches=False):
+    """
+    Batch a function such that each batch is smaller than a threshold.
+
+    Parameters
+    ----------
+    func : callable
+        A jittable function with positional arguments only, with inputs and
+        outputs pytrees of arrays.
+    max_io_nbytes : int
+        The maximum number of input + output bytes in each batch.
+    in_axes : pytree of ints, default 0
+        A tree matching the structure of the function input, indicating along
+        which axes each array should be batched. If a single integer, it is
+        used for all arrays.
+    out_axes : pytree of ints, default 0
+        The same for outputs.
+    return_nbatches : bool, default False
+        If True, the number of batches is returned as a second output.
+
+    Returns
+    -------
+    batched_func : callable
+        A function with the same signature as `func`, but that processes the
+        input and output in batches in a loop.
+    """
+
+    def expand_axes(axes, tree):
+        if isinstance(axes, int):
+            return tree_util.tree_map(lambda _: axes, tree)
+        return tree_util.tree_map(lambda _, axis: axis, tree, axes)
+
+    def extract_size(axes, tree):
+        sizes = tree_util.tree_map(lambda x, axis: x.shape[axis], tree, axes)
+        sizes, _ = tree_util.tree_flatten(sizes)
+        assert all(s == sizes[0] for s in sizes)
+        return sizes[0]
+
+    def sum_nbytes(tree):
+        def nbytes(x):
+            return math.prod(x.shape) * x.dtype.itemsize
+        return tree_util.tree_reduce(lambda size, x: size + nbytes(x), tree, 0)
+
+    def next_divisor_small(dividend, min_divisor):
+        for divisor in range(min_divisor, int(math.sqrt(dividend)) + 1):
+            if dividend % divisor == 0:
+                return divisor
+        return dividend
+
+    def next_divisor_large(dividend, min_divisor):
+        max_inv_divisor = dividend // min_divisor
+        for inv_divisor in range(max_inv_divisor, 0, -1):
+            if dividend % inv_divisor == 0:
+                return dividend // inv_divisor
+        return dividend
+
+    def next_divisor(dividend, min_divisor):
+        if min_divisor * min_divisor <= dividend:
+            return next_divisor_small(dividend, min_divisor)
+        return next_divisor_large(dividend, min_divisor)
+
+    def move_axes_out(axes, tree):
+        def move_axis_out(axis, x):
+            if axis != 0:
+                return jnp.moveaxis(x, axis, 0)
+            return x
+        return tree_util.tree_map(move_axis_out, axes, tree)
+
+    def move_axes_in(axes, tree):
+        def move_axis_in(axis, x):
+            if axis != 0:
+                return jnp.moveaxis(x, 0, axis)
+            return x
+        return tree_util.tree_map(move_axis_in, axes, tree)
+
+    def batch(tree, nbatches):
+        def batch(x):
+            return x.reshape((nbatches, x.shape[0] // nbatches) + x.shape[1:])
+        return tree_util.tree_map(batch, tree)
+
+    def unbatch(tree):
+        def unbatch(x):
+            return x.reshape((x.shape[0] * x.shape[1],) + x.shape[2:])
+        return tree_util.tree_map(unbatch, tree)
+
+    def check_same(tree1, tree2):
+        def check_same(x1, x2):
+            assert x1.shape == x2.shape
+            assert x1.dtype == x2.dtype
+        tree_util.tree_map(check_same, tree1, tree2)
+
+    initial_in_axes = in_axes
+    initial_out_axes = out_axes
+
+    @jax.jit
+    @functools.wraps(func)
+    def batched_func(*args):
+        example_result = jax.eval_shape(func, *args)
+
+        in_axes = expand_axes(initial_in_axes, args)
+        out_axes = expand_axes(initial_out_axes, example_result)
+
+        in_size = extract_size(in_axes, args)
+        out_size = extract_size(out_axes, example_result)
+        assert in_size == out_size
+        size = in_size
+
+        total_nbytes = sum_nbytes(args) + sum_nbytes(example_result)
+        min_nbatches = total_nbytes // max_io_nbytes + bool(total_nbytes % max_io_nbytes)
+        nbatches = next_divisor(size, min_nbatches)
+        assert 1 <= nbatches <= size
+        assert size % nbatches == 0
+        assert total_nbytes % nbatches == 0
+
+        batch_nbytes = total_nbytes // nbatches
+        if batch_nbytes > max_io_nbytes:
+            assert size == nbatches
+            warnings.warn(f'batch_nbytes = {batch_nbytes} > max_io_nbytes = {max_io_nbytes}')
+
+        def loop(_, args):
+            args = move_axes_in(in_axes, args)
+            result = func(*args)
+            result = move_axes_out(out_axes, result)
+            return None, result
+
+        args = move_axes_out(in_axes, args)
+        args = batch(args, nbatches)
+        _, result = lax.scan(loop, None, args)
+        result = unbatch(result)
+        result = move_axes_in(out_axes, result)
+
+        check_same(example_result, result)
+
+        if return_nbatches:
+            return result, nbatches
+        return result
+
+    return batched_func
+
+@tree_util.register_pytree_node_class
+class LeafDict(dict):
+    """ dictionary that acts as a leaf in jax pytrees, to store compile-time
+    values """
+
+    def tree_flatten(self):
+        return (), self
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        return aux_data
+
+    def __repr__(self):
+        return f'{__class__.__name__}({super().__repr__()})'

bartz/mcmcloop.py

@@ -52,7 +52,7 @@ def run_mcmc(bart, n_burn, n_save, n_skip, callback, key):
     n_save : int
         The number of iterations to save.
     n_skip : int
-        The number of iterations to skip between each saved iteration.
+        The number of iterations to skip between each saved iteration, plus 1.
     callback : callable
         An arbitrary function run at each iteration, called with the following
         arguments, passed by keyword:
@@ -105,16 +105,19 @@ def run_mcmc(bart, n_burn, n_save, n_skip, callback, key):
         output = {key: bart[key] for key in tracelist}
         return (bart, i_total + 1, i_skip + 1, key), output
 
+    def empty_trace(bart, tracelist):
+        return {
+            key: jnp.empty((0,) + bart[key].shape, bart[key].dtype)
+            for key in tracelist
+        }
+
     if n_burn > 0:
         carry = bart, 0, 0, key
         burnin_loop = functools.partial(inner_loop, tracelist=tracelist_burnin, burnin=True)
         (bart, i_total, _, key), burnin_trace = lax.scan(burnin_loop, carry, None, n_burn)
     else:
         i_total = 0
-        burnin_trace = {
-            key: jnp.empty((0,) + bart[key].shape, bart[key].dtype)
-            for key in tracelist_burnin
-        }
+        burnin_trace = empty_trace(bart, tracelist_burnin)
 
     def outer_loop(carry, _):
         bart, i_total, key = carry
@@ -124,8 +127,11 @@ def run_mcmc(bart, n_burn, n_save, n_skip, callback, key):
         output = {key: bart[key] for key in tracelist_main}
         return (bart, i_total, key), output
 
-    carry = bart, i_total, key
-    (bart, _, _), main_trace = lax.scan(outer_loop, carry, None, n_save)
+    if n_save > 0:
+        carry = bart, i_total, key
+        (bart, _, _), main_trace = lax.scan(outer_loop, carry, None, n_save)
+    else:
+        main_trace = empty_trace(bart, tracelist_main)
 
     return bart, burnin_trace, main_trace
 
@@ -133,7 +139,8 @@ def run_mcmc(bart, n_burn, n_save, n_skip, callback, key):
 
 @functools.lru_cache
     # cache to make the callback function object unique, such that the jit
-    # of run_mcmc recognizes it
+    # of run_mcmc recognizes it => with the callback state, I can make
+    # printevery a runtime quantity
 def make_simple_print_callback(printevery):
     """
     Create a logging callback function for MCMC iterations.
@@ -155,11 +162,12 @@ def make_simple_print_callback(printevery):
         grow_acc = bart['grow_acc_count'] / bart['grow_prop_count']
         prune_acc = bart['prune_acc_count'] / bart['prune_prop_count']
         n_total = n_burn + n_save * n_skip
-        debug.callback(simple_print_callback_impl, burnin, i_total, n_total, grow_prop, grow_acc, prune_prop, prune_acc, printevery)
+        printcond = (i_total + 1) % printevery == 0
+        debug.callback(_simple_print_callback, burnin, i_total, n_total, grow_prop, grow_acc, prune_prop, prune_acc, printcond)
     return callback
 
-def simple_print_callback_impl(burnin, i_total, n_total, grow_prop, grow_acc, prune_prop, prune_acc, printevery):
-    if (i_total + 1) % printevery == 0:
+def _simple_print_callback(burnin, i_total, n_total, grow_prop, grow_acc, prune_prop, prune_acc, printcond):
+    if printcond:
         burnin_flag = ' (burnin)' if burnin else ''
         total_str = str(n_total)
         ndigits = len(total_str)

bartz/mcmcstep.py

@@ -34,7 +34,6 @@ range of possible values.
 """
 
 import functools
-import math
 
 import jax
 from jax import random
@@ -55,6 +54,7 @@ def init(*,
     small_float=jnp.float32,
     large_float=jnp.float32,
     min_points_per_leaf=None,
+    suffstat_batch_size='auto',
     ):
     """
     Make a BART posterior sampling MCMC initial state.
@@ -82,6 +82,10 @@ 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.
 
     Returns
     -------
@@ -104,8 +108,9 @@ def init(*,
             The number of grow/prune proposals made during one full MCMC cycle.
         'grow_acc_count', 'prune_acc_count' : int
             The number of grow/prune moves accepted during one full MCMC cycle.
-        'p_nonterminal' : large_float array (d - 1,)
-            The probability of a nonterminal node at each depth.
+        'p_nonterminal' : large_float array (d,)
+            The probability of a nonterminal node at each depth, padded with a
+            zero.
         'sigma2_alpha' : large_float
             The shape parameter of the inverse gamma prior on the noise variance.
         'sigma2_beta' : large_float
@@ -121,18 +126,36 @@ def init(*,
         'affluence_trees' : bool array (num_trees, 2 ** (d - 1)) or None
             Whether a non-bottom leaf nodes contains twice `min_points_per_leaf`
             datapoints. If `min_points_per_leaf` is not specified, this is None.
+        '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
+                The dtype for scalars, small arrays, and arrays which require
+                accuracy.
+            'require_min_points' : bool
+                Whether the `min_points_per_leaf` parameter is specified.
     """
 
     p_nonterminal = jnp.asarray(p_nonterminal, large_float)
-    max_depth = p_nonterminal.size + 1
+    p_nonterminal = jnp.pad(p_nonterminal, (0, 1))
+    max_depth = p_nonterminal.size
 
     @functools.partial(jax.vmap, in_axes=None, out_axes=0, axis_size=num_trees)
     def make_forest(max_depth, dtype):
         return grove.make_tree(max_depth, dtype)
 
+    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)
+
     bart = dict(
         leaf_trees=make_forest(max_depth, small_float),
-        var_trees=make_forest(max_depth - 1, grove.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=jnp.ones((), large_float),
@@ -143,9 +166,9 @@ def init(*,
         p_nonterminal=p_nonterminal,
         sigma2_alpha=jnp.asarray(sigma2_alpha, large_float),
         sigma2_beta=jnp.asarray(sigma2_beta, large_float),
-        max_split=max_split,
-        y=jnp.asarray(y, small_float),
-        X=X,
+        max_split=jnp.asarray(max_split),
+        y=y,
+        X=jnp.asarray(X),
         min_points_per_leaf=(
             None if min_points_per_leaf is None else
             jnp.asarray(min_points_per_leaf)
@@ -154,10 +177,39 @@ def init(*,
             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(
+            suffstat_batch_size=suffstat_batch_size,
+            small_float=small_float,
+            large_float=large_float,
+            require_min_points=min_points_per_leaf is not None,
+        ),
     )
 
     return bart
 
+def _choose_suffstat_batch_size(size, y):
+    if size == 'auto':
+        try:
+            device = y.devices().pop()
+        except jax.errors.ConcretizationTypeError:
+            device = jax.devices()[0]
+        platform = device.platform
+
+        if platform == 'cpu':
+            return None
+                # maybe I should batch residuals (not counts) for numerical
+                # accuracy, even if it's slower
+        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
+
 def step(bart, key):
     """
     Perform one full MCMC step on a BART state.
@@ -196,11 +248,14 @@ def sample_trees(bart, key):
 
     Notes
     -----
-    This function zeroes the proposal counters.
+    This function zeroes the proposal counters before using them.
     """
+    bart = bart.copy()
     key, subkey = random.split(key)
     grow_moves, prune_moves = sample_moves(bart, subkey)
-    return accept_moves_and_sample_leaves(bart, grow_moves, prune_moves, key)
+    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)
 
 def sample_moves(bart, key):
     """
@@ -216,20 +271,7 @@ def sample_moves(bart, key):
     Returns
     -------
     grow_moves, prune_moves : dict
-        The proposals for grow and prune moves, with these fields:
-
-        'allowed' : bool array (num_trees,)
-            Whether the move is possible.
-        'node' : int array (num_trees,)
-            The index of the leaf to grow or node to prune.
-        'var_tree' : int array (num_trees, 2 ** (d - 1),)
-            The new decision axes of the tree.
-        'split_tree' : int array (num_trees, 2 ** (d - 1),)
-            The new decision boundaries of the tree.
-        '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. For the prune move, the ratio is inverted.
+        The proposals for grow and prune moves. See `grow_move` and `prune_move`.
     """
     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)
@@ -260,7 +302,7 @@ def grow_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, ke
         Whether a leaf has enough points to be grown.
     max_split : array (p,)
         The maximum split index for each variable.
-    p_nonterminal : array (d - 1,)
+    p_nonterminal : array (d,)
         The probability of a nonterminal node at each depth.
     key : jax.dtypes.prng_key array
         A jax random key.
@@ -292,16 +334,16 @@ def grow_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, ke
     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)
-    new_split_tree = split_tree.at[leaf_to_grow].set(split.astype(split_tree.dtype))
+    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, new_split_tree)
+    ratio = compute_partial_ratio(num_growable, num_prunable, p_nonterminal, leaf_to_grow, split_tree)
 
     return dict(
         allowed=allowed,
         node=leaf_to_grow,
-        var_tree=var_tree,
-        split_tree=new_split_tree,
         partial_ratio=ratio,
+        var_tree=var_tree,
+        split_tree=split_tree,
     )
 
 def choose_leaf(split_tree, affluence_tree, key):
@@ -464,7 +506,7 @@ 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, grove.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
@@ -569,7 +611,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, initial_split_tree, new_split_tree):
+def compute_partial_ratio(num_growable, num_prunable, p_nonterminal, leaf_to_grow, new_split_tree):
     """
     Compute the product of the transition and prior ratios of a grow move.
 
@@ -580,12 +622,10 @@ def compute_partial_ratio(num_growable, num_prunable, p_nonterminal, leaf_to_gro
     num_prunable : int
         The number of leaf parents that could be pruned, after converting the
         leaf to be grown to a non-terminal node.
-    p_nonterminal : array (d - 1,)
+    p_nonterminal : array (d,)
         The probability of a nonterminal node at each depth.
     leaf_to_grow : int
         The index of the leaf to grow.
-    initial_split_tree : array (2 ** (d - 1),)
-        The splitting points of the tree, before the leaf is grown.
     new_split_tree : array (2 ** (d - 1),)
         The splitting points of the tree, after the leaf is grown.
 
@@ -600,14 +640,18 @@ 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
 
-    prune_was_allowed = prune_allowed(initial_split_tree)
-    p_grow = jnp.where(prune_was_allowed, 0.5, 1)
+    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
+
+    p_grow = jnp.where(prune_allowed, 0.5, 1)
 
     trans_ratio = num_growable / (p_grow * num_prunable)
 
-    depth = grove.tree_depths(initial_split_tree.size)[leaf_to_grow]
+    depth = grove.tree_depths(new_split_tree.size)[leaf_to_grow]
     p_parent = p_nonterminal[depth]
-    cp_children = 1 - p_nonterminal.at[depth + 1].get(mode='fill', fill_value=0)
+    cp_children = 1 - p_nonterminal[depth + 1]
     tree_ratio = cp_children * cp_children * p_parent / (1 - p_parent)
 
     return trans_ratio * tree_ratio
@@ -626,7 +670,7 @@ def prune_move(var_tree, split_tree, affluence_tree, max_split, p_nonterminal, k
         Whether a leaf has enough points to be grown.
     max_split : array (p,)
         The maximum split index for each variable.
-    p_nonterminal : array (d - 1,)
+    p_nonterminal : array (d,)
         The probability of a nonterminal node at each depth.
     key : jax.dtypes.prng_key array
         A jax random key.
@@ -639,28 +683,20 @@ 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 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 node to prune.
         '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)
-    allowed = prune_allowed(split_tree)
-
-    new_split_tree = split_tree.at[node_to_prune].set(0)
+    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, new_split_tree, split_tree)
+    ratio = compute_partial_ratio(num_growable, num_prunable, p_nonterminal, node_to_prune, split_tree)
 
     return dict(
         allowed=allowed,
         node=node_to_prune,
-        var_tree=var_tree,
-        split_tree=new_split_tree,
         partial_ratio=ratio, # it is inverted in accept_move_and_sample_leaves
     )
 
@@ -702,29 +738,37 @@ def choose_leaf_parent(split_tree, affluence_tree, key):
 
     return node_to_prune, num_prunable, num_growable
 
-def prune_allowed(split_tree):
+def accept_moves_and_sample_leaves(bart, grow_moves, prune_moves, grow_leaf_indices, key):
     """
-    Return whether a prune move is allowed.
+    Accept or reject the proposed moves and sample the new leaf values.
 
     Parameters
     ----------
-    split_tree : array (2 ** (d - 1),)
-        The splitting points of the tree.
+    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.
+    key : jax.dtypes.prng_key array
+        A jax random key.
 
     Returns
     -------
-    allowed : bool
-        Whether a prune move is allowed.
+    bart : dict
+        The new BART mcmc state.
     """
-    return split_tree.at[1].get(mode='fill', fill_value=0).astype(bool)
-
-def accept_moves_and_sample_leaves(bart, grow_moves, prune_moves, key):
     bart = bart.copy()
     def loop(carry, item):
         resid = carry.pop('resid')
         resid, carry, trees = accept_move_and_sample_leaves(
             bart['X'],
             len(bart['leaf_trees']),
+            bart['opt']['suffstat_batch_size'],
             resid,
             bart['sigma2'],
             bart['min_points_per_leaf'],
@@ -740,11 +784,11 @@ def accept_moves_and_sample_leaves(bart, grow_moves, prune_moves, key):
     carry['resid'] = bart['resid']
     items = (
         bart['leaf_trees'],
-        bart['var_trees'],
         bart['split_trees'],
         bart['affluence_trees'],
         grow_moves,
         prune_moves,
+        grow_leaf_indices,
         random.split(key, len(bart['leaf_trees'])),
     )
     carry, trees = lax.scan(loop, carry, items)
@@ -752,11 +796,50 @@ def accept_moves_and_sample_leaves(bart, grow_moves, prune_moves, key):
     bart.update(trees)
     return bart
 
-def accept_move_and_sample_leaves(X, ntree, resid, sigma2, min_points_per_leaf, counts, leaf_tree, var_tree, split_tree, affluence_tree, grow_move, prune_move, key):
-    
-    # compute leaf indices according to grow move tree
-    traverse_tree = jax.vmap(grove.traverse_tree, in_axes=(1, None, None))
-    grow_leaf_indices = traverse_tree(X, grow_move['var_tree'], grow_move['split_tree'])
+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):
+    """
+    Accept or reject a proposed move and sample the new leaf values.
+
+    Parameters
+    ----------
+    X : int array (p, n)
+        The predictors.
+    ntree : int
+        The number of trees in the forest.
+    suffstat_batch_size : int, None
+        The batch size for computing sufficient statistics.
+    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.
+    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.
+
+    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']
@@ -782,10 +865,7 @@ def accept_move_and_sample_leaves(X, ntree, resid, sigma2, min_points_per_leaf,
     resid += leaf_tree[leaf_indices]
 
     # aggregate residuals and count units per leaf
-    grow_resid_tree = jnp.zeros_like(leaf_tree, sigma2.dtype)
-    grow_resid_tree = grow_resid_tree.at[grow_leaf_indices].add(resid)
-    grow_count_tree = jnp.zeros_like(leaf_tree, grove.minimal_unsigned_dtype(resid.size))
-    grow_count_tree = grow_count_tree.at[grow_leaf_indices].add(1)
+    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
@@ -833,10 +913,10 @@ def accept_move_and_sample_leaves(X, ntree, resid, sigma2, min_points_per_leaf,
 
     # pick trees for chosen move
     trees = {}
-    var_tree = jnp.where(do_grow, grow_move['var_tree'], var_tree)
     split_tree = jnp.where(do_grow, grow_move['split_tree'], split_tree)
-    var_tree = jnp.where(do_prune, prune_move['var_tree'], var_tree)
-    split_tree = jnp.where(do_prune, prune_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]))
     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)
@@ -869,13 +949,60 @@ def accept_move_and_sample_leaves(X, ntree, resid, sigma2, min_points_per_leaf,
     # pack trees
     trees = {
         'leaf_trees': leaf_tree,
-        'var_trees': var_tree,
         'split_trees': split_tree,
         'affluence_trees': affluence_tree,
     }
 
     return resid, counts, trees
 
+def sufficient_stat(resid, leaf_indices, tree_size, batch_size):
+    """
+    Compute the sufficient statistics for the likelihood ratio of a tree move.
+
+    Parameters
+    ----------
+    resid : float array (n,)
+        The residuals (data minus forest value).
+    leaf_indices : int array (n,)
+        The leaf indices of the tree (in which leaf each data point falls into).
+    tree_size : int
+        The size of the tree array (2 ** d).
+    batch_size : int, None
+        The 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
+
+def _aggregate_scatter(values, indices, size, dtype):
+    return (jnp
+        .zeros(size, dtype)
+        .at[indices]
+        .add(values)
+    )
+
+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):
     """
     Compute the probability of proposing a prune move after doing a grow move.

bartz/prepcovars.py

@@ -27,8 +27,10 @@ import functools
 import jax
 from jax import numpy as jnp
 
+from . import jaxext
 from . import grove
 
+@functools.partial(jax.jit, static_argnums=(1,))
 def quantilized_splits_from_matrix(X, max_bins):
     """
     Determine bins that make the distribution of each predictor uniform.
@@ -52,48 +54,41 @@ def quantilized_splits_from_matrix(X, max_bins):
         The number of actually used values in each row of `splits`.
     """
     out_length = min(max_bins, X.shape[1]) - 1
-    return quantilized_splits_from_matrix_impl(X, out_length)
+    # return _quantilized_splits_from_matrix(X, out_length)
+    @functools.partial(jaxext.autobatch, max_io_nbytes=500_000_000)
+    def func(X):
+        return _quantilized_splits_from_matrix(X, out_length)
+    return func(X)
 
 @functools.partial(jax.vmap, in_axes=(0, None))
-def quantilized_splits_from_matrix_impl(x, out_length):
-    huge = huge_value(x)
-    u = jnp.unique(x, size=x.size, fill_value=huge)
-    actual_length = jnp.count_nonzero(u < huge) - 1
-    midpoints = (u[1:] + u[:-1]) / 2
+def _quantilized_splits_from_matrix(x, out_length):
+    huge = jaxext.huge_value(x)
+    u, actual_length = jaxext.unique(x, size=x.size, fill_value=huge)
+    actual_length -= 1
+    if jnp.issubdtype(x.dtype, jnp.integer):
+        midpoints = u[:-1] + jaxext.ensure_unsigned(u[1:] - u[:-1]) // 2
+        indices = jnp.arange(midpoints.size, dtype=jaxext.minimal_unsigned_dtype(midpoints.size - 1))
+        midpoints = jnp.where(indices < actual_length, midpoints, huge)
+    else:
+        midpoints = (u[1:] + u[:-1]) / 2
     indices = jnp.linspace(-1, actual_length, out_length + 2)[1:-1]
-    indices = jnp.around(indices).astype(grove.minimal_unsigned_dtype(midpoints.size - 1))
+    indices = jnp.around(indices).astype(jaxext.minimal_unsigned_dtype(midpoints.size - 1))
         # indices calculation with float rather than int to avoid potential
         # overflow with int32, and to round to nearest instead of rounding down
     decimated_midpoints = midpoints[indices]
     truncated_midpoints = midpoints[:out_length]
     splits = jnp.where(actual_length > out_length, decimated_midpoints, truncated_midpoints)
     max_split = jnp.minimum(actual_length, out_length)
-    max_split = max_split.astype(grove.minimal_unsigned_dtype(out_length))
+    max_split = max_split.astype(jaxext.minimal_unsigned_dtype(out_length))
     return splits, max_split
 
-def huge_value(x):
-    """
-    Return the maximum value that can be stored in `x`.
-
-    Parameters
-    ----------
-    x : array
-        A numerical numpy or jax array.
-
-    Returns
-    -------
-    maxval : scalar
-        The maximum value allowed by `x`'s type (+inf for floats).
-    """
-    if jnp.issubdtype(x.dtype, jnp.integer):
-        return jnp.iinfo(x.dtype).max
-    else:
-        return jnp.inf
-
+@jax.jit
 def bin_predictors(X, splits):
     """
     Bin the predictors according to the given splits.
 
+    A value ``x`` is mapped to bin ``i`` iff ``splits[i - 1] < x <= splits[i]``.
+
     Parameters
     ----------
     X : array (p, n)
@@ -110,9 +105,9 @@ def bin_predictors(X, splits):
         A matrix with `p` predictors and `n` observations, where each predictor
         has been replaced by the index of the bin it falls into.
     """
-    return bin_predictors_impl(X, splits)
+    return _bin_predictors(X, splits)
 
 @jax.vmap
-def bin_predictors_impl(x, splits):
-    dtype = grove.minimal_unsigned_dtype(splits.size)
+def _bin_predictors(x, splits):
+    dtype = jaxext.minimal_unsigned_dtype(splits.size)
     return jnp.searchsorted(splits, x).astype(dtype)

{bartz-0.1.0.dist-info → bartz-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bartz
-Version: 0.1.0
+Version: 0.2.1
 Summary: A JAX implementation of BART
 Home-page: https://github.com/Gattocrucco/bartz
 License: MIT
@@ -20,7 +20,13 @@ Project-URL: Bug Tracker, https://github.com/Gattocrucco/bartz/issues
 Project-URL: Repository, https://github.com/Gattocrucco/bartz
 Description-Content-Type: text/markdown
 
+[![PyPI](https://img.shields.io/pypi/v/bartz)](https://pypi.org/project/bartz/)
+
 # BART vectoriZed
 
 A branchless vectorized implementation of Bayesian Additive Regression Trees (BART) in JAX.
 
+BART is a nonparametric Bayesian regression technique. Given predictors $X$ and responses $y$, BART finds a function to predict $y$ given $X$. The result of the inference is a sample of possible functions, representing the uncertainty over the determination of the function.
+
+This Python module provides an implementation of BART that runs on GPU, to process large datasets faster. It is also a good on CPU. Most other implementations of BART are for R, and run on CPU only.
+

bartz-0.2.1.dist-info/RECORD

@@ -0,0 +1,13 @@
+bartz/BART.py,sha256=pRG7mALenknX2JHqY-VyhO9-evDgEC6hWBp4jpecBdM,15801
+bartz/__init__.py,sha256=E96vsP0bZ8brejpZmEmRoXuMsUdinO_B_SKUUl1rLsg,1448
+bartz/_version.py,sha256=PmcQ2PI2oP8irnLtJLJby2YfW6sBvLAmL-VpABzTqwc,22
+bartz/debug.py,sha256=9ZH-JfwZVu5OPhHBEyXQHAU5H9KIu1vxLK7yNv4m4Ew,5314
+bartz/grove.py,sha256=Wj_7jHl9w3uwuVdH4hoeXowimGpdRE2lGIzr4aDkzsI,8291
+bartz/jaxext.py,sha256=VYA41D5F7DYcAAVtkcZtEN927HxQGOOQM-uGsgr2CPc,10996
+bartz/mcmcloop.py,sha256=lheLrjVxmlyQzc_92zeNsFhdkrhEWQEjoAWFbVzknnw,7701
+bartz/mcmcstep.py,sha256=6fzNMumXjMe6Fj6zoHLTf1D42JuAiQyGHfr6l1Bwrnk,39450
+bartz/prepcovars.py,sha256=iiQ0WjSj4--l5DgPW626Qg2SSB6ljnaaUsBz_A8kFrI,4634
+bartz-0.2.1.dist-info/LICENSE,sha256=heuIJZQK9IexJYC-fYHoLUrgj8HG8yS3G072EvKh-94,1073
+bartz-0.2.1.dist-info/METADATA,sha256=eGxicC1iR-Bpjk1uKn50g6FxdFfq9S70nl7m5GmXO14,1490
+bartz-0.2.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+bartz-0.2.1.dist-info/RECORD,,

bartz-0.1.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-bartz/__init__.py,sha256=40tX5XHoTiGnZcoeogVpyNOM_5rbHt-Y6zTI0NS7OA4,1345
-bartz/_version.py,sha256=IMjkMO3twhQzluVTo8Z6rE7Eg-9U79_LGKMcsWLKBkY,22
-bartz/debug.py,sha256=_HOjDieipAgliP6B6C0UMgz-mVgmeZ3zmtzVe-iMGtY,5289
-bartz/grove.py,sha256=LHhnvNKLb-jxUf4YjP927Hf9txkXynhMZ2ejtMRWZl4,8353
-bartz/interface.py,sha256=INyNuHzFySwXAsXofVZDpTsMv78AR_3VCvAHbZFh92c,15724
-bartz/jaxext.py,sha256=FK5j1zfW1yR4-yPKcD7ZvKSkVQ5--jHjQpVCl4n4gXY,2844
-bartz/mcmcloop.py,sha256=xTxC1AkNX8jCrMArblvlMjnjMh80q1M3a6ZGrDdfsFI,7423
-bartz/mcmcstep.py,sha256=6zkpTqgIrapeVy9mhy6BlsIO0s26HwBRDfw_6dVMmZA,35207
-bartz/prepcovars.py,sha256=3ddDOtNNop3Ba2Kgy_dZ6apFydtwaEXH3uXSmmKf9Fs,4421
-bartz-0.1.0.dist-info/LICENSE,sha256=heuIJZQK9IexJYC-fYHoLUrgj8HG8yS3G072EvKh-94,1073
-bartz-0.1.0.dist-info/METADATA,sha256=8YYlbCf7frDtT2of6tNlnBbuGqyO8YyYlED8OXSiBpA,933
-bartz-0.1.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-bartz-0.1.0.dist-info/RECORD,,