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

bartz/prepcovars.py

@@ -22,64 +22,113 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-import functools
+"""Functions to preprocess data."""
 
-import jax
+from functools import partial
+
+from jax import jit, vmap
 from jax import numpy as jnp
+from jaxtyping import Array, Float, Integer, Real, UInt
+
+from bartz.jaxext import autobatch, minimal_unsigned_dtype, unique
+
 
-from . import jaxext
+def parse_xinfo(
+    xinfo: Float[Array, 'p m'],
+) -> tuple[Float[Array, 'p m'], UInt[Array, ' p']]:
+    """Parse pre-defined splits in the format of the R package BART.
+
+    Parameters
+    ----------
+    xinfo
+        A matrix with the cutpoins to use to bin each predictor. Each row shall
+        contain a sorted list of cutpoints for a predictor. If there are less
+        cutpoints than the number of columns in the matrix, fill the remaining
+        cells with NaN.
+
+        `xinfo` shall be a matrix even if `x_train` is a dataframe.
+
+    Returns
+    -------
+    splits : Float[Array, 'p m']
+        `xinfo` modified by replacing nan with a large value.
+    max_split : UInt[Array, 'p']
+        The number of non-nan elements in each row of `xinfo`.
+    """
+    is_not_nan = ~jnp.isnan(xinfo)
+    max_split = jnp.sum(is_not_nan, axis=1)
+    max_split = max_split.astype(minimal_unsigned_dtype(xinfo.shape[1]))
+    huge = _huge_value(xinfo)
+    splits = jnp.where(is_not_nan, xinfo, huge)
+    return splits, max_split
 
 
-@functools.partial(jax.jit, static_argnums=(1,))
-def quantilized_splits_from_matrix(X, max_bins):
+@partial(jit, static_argnums=(1,))
+def quantilized_splits_from_matrix(
+    X: Real[Array, 'p n'], max_bins: int
+) -> tuple[Real[Array, 'p m'], UInt[Array, ' p']]:
     """
     Determine bins that make the distribution of each predictor uniform.
 
     Parameters
     ----------
-    X : array (p, n)
+    X
         A matrix with `p` predictors and `n` observations.
-    max_bins : int
+    max_bins
         The maximum number of bins to produce.
 
     Returns
     -------
-    splits : array (p, m)
+    splits : Real[Array, 'p m']
         A matrix containing, for each predictor, the boundaries between bins.
         `m` is ``min(max_bins, n) - 1``, which is an upper bound on the number
         of splits. Each predictor may have a different number of splits; unused
         values at the end of each row are filled with the maximum value
         representable in the type of `X`.
-    max_split : array (p,)
+    max_split : UInt[Array, ' p']
         The number of actually used values in each row of `splits`.
+
+    Raises
+    ------
+    ValueError
+        If `X` has no columns or if `max_bins` is less than 1.
     """
     out_length = min(max_bins, X.shape[1]) - 1
 
-    # return _quantilized_splits_from_matrix(X, out_length)
-    @functools.partial(jaxext.autobatch, max_io_nbytes=2**29)
+    if out_length < 0:
+        msg = f'{X.shape[1]=} and {max_bins=}, they should be both at least 1.'
+        raise ValueError(msg)
+
+    @partial(autobatch, max_io_nbytes=2**29)
     def quantilize(X):
+        # wrap this function because autobatch needs traceable args
         return _quantilized_splits_from_matrix(X, out_length)
 
     return quantilize(X)
 
 
-@functools.partial(jax.vmap, in_axes=(0, None))
-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
+@partial(vmap, in_axes=(0, None))
+def _quantilized_splits_from_matrix(
+    x: Real[Array, 'p n'], out_length: int
+) -> tuple[Real[Array, 'p m'], UInt[Array, ' p']]:
+    # find the sorted unique values in x
+    huge = _huge_value(x)
+    u, actual_length = unique(x, size=x.size, fill_value=huge)
+
+    # compute the midpoints between each unique value
     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)
+        midpoints = u[:-1] + _ensure_unsigned(u[1:] - u[:-1]) // 2
     else:
-        midpoints = (u[1:] + u[:-1]) / 2
+        midpoints = u[:-1] + (u[1:] - u[:-1]) / 2
+        # using x_i + (x_i+1 - x_i) / 2 instead of (x_i + x_i+1) / 2 is to
+        # avoid overflow
+    actual_length -= 1
+    if midpoints.size:
+        midpoints = midpoints.at[actual_length].set(huge)
+
+    # take a subset of the midpoints if there are more than the requested maximum
     indices = jnp.linspace(-1, actual_length, out_length + 2)[1:-1]
-    indices = jnp.around(indices).astype(
-        jaxext.minimal_unsigned_dtype(midpoints.size - 1)
-    )
+    indices = jnp.around(indices).astype(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]
@@ -88,41 +137,92 @@ def _quantilized_splits_from_matrix(x, out_length):
         actual_length > out_length, decimated_midpoints, truncated_midpoints
     )
     max_split = jnp.minimum(actual_length, out_length)
-    max_split = max_split.astype(jaxext.minimal_unsigned_dtype(out_length))
+    max_split = max_split.astype(minimal_unsigned_dtype(out_length))
     return splits, max_split
 
 
-@functools.partial(jax.jit, static_argnums=(1,))
-def uniform_splits_from_matrix(X, num_bins):
+def _huge_value(x: Array) -> int | float:
+    """
+    Return the maximum value that can be stored in `x`.
+
+    Parameters
+    ----------
+    x
+        A numerical numpy or jax array.
+
+    Returns
+    -------
+    The maximum value allowed by `x`'s type (finite for floats).
+    """
+    if jnp.issubdtype(x.dtype, jnp.integer):
+        return jnp.iinfo(x.dtype).max
+    else:
+        return float(jnp.finfo(x.dtype).max)
+
+
+def _ensure_unsigned(x: Integer[Array, '*shape']) -> UInt[Array, '*shape']:
+    """If x has signed integer type, cast it to the unsigned dtype of the same size."""
+    return x.astype(_signed_to_unsigned(x.dtype))
+
+
+def _signed_to_unsigned(int_dtype: jnp.dtype) -> jnp.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
+    match int_dtype:
+        case jnp.int8:
+            return jnp.uint8
+        case jnp.int16:
+            return jnp.uint16
+        case jnp.int32:
+            return jnp.uint32
+        case jnp.int64:
+            return jnp.uint64
+        case _:
+            msg = f'unexpected integer type {int_dtype}'
+            raise TypeError(msg)
+
+
+@partial(jit, static_argnums=(1,))
+def uniform_splits_from_matrix(
+    X: Real[Array, 'p n'], num_bins: int
+) -> tuple[Real[Array, 'p m'], UInt[Array, ' p']]:
     """
     Make an evenly spaced binning grid.
 
     Parameters
     ----------
-    X : array (p, n)
+    X
         A matrix with `p` predictors and `n` observations.
-    num_bins : int
+    num_bins
         The number of bins to produce.
 
     Returns
     -------
-    splits : array (p, num_bins - 1)
+    splits : Real[Array, 'p m']
         A matrix containing, for each predictor, the boundaries between bins.
         The excluded endpoints are the minimum and maximum value in each row of
         `X`.
-    max_split : array (p,)
+    max_split : UInt[Array, ' p']
         The number of cutpoints in each row of `splits`, i.e., ``num_bins - 1``.
     """
     low = jnp.min(X, axis=1)
     high = jnp.max(X, axis=1)
     splits = jnp.linspace(low, high, num_bins + 1, axis=1)[:, 1:-1]
     assert splits.shape == (X.shape[0], num_bins - 1)
-    max_split = jnp.full(*splits.shape, jaxext.minimal_unsigned_dtype(num_bins - 1))
+    max_split = jnp.full(*splits.shape, minimal_unsigned_dtype(num_bins - 1))
     return splits, max_split
 
 
-@functools.partial(jax.jit, static_argnames=('method',))
-def bin_predictors(X, splits, **kw):
+@partial(jit, static_argnames=('method',))
+def bin_predictors(
+    X: Real[Array, 'p n'], splits: Real[Array, 'p m'], **kw
+) -> UInt[Array, 'p n']:
     """
     Bin the predictors according to the given splits.
 
@@ -130,27 +230,25 @@ def bin_predictors(X, splits, **kw):
 
     Parameters
     ----------
-    X : array (p, n)
+    X
         A matrix with `p` predictors and `n` observations.
-    splits : array (p, m)
+    splits
         A matrix containing, for each predictor, the boundaries between bins.
         `m` is the maximum number of splits; each row may have shorter
         actual length, marked by padding unused locations at the end of the
         row with the maximum value allowed by the type.
-    **kw : dict
+    **kw
         Additional arguments are passed to `jax.numpy.searchsorted`.
 
     Returns
     -------
-    X_binned : int array (p, n)
-        A matrix with `p` predictors and `n` observations, where each predictor
-        has been replaced by the index of the bin it falls into.
+    `X` but with each value replaced by the index of the bin it falls into.
     """
 
-    @functools.partial(jaxext.autobatch, max_io_nbytes=2**29)
-    @jax.vmap
+    @partial(autobatch, max_io_nbytes=2**29)
+    @vmap
     def bin_predictors(x, splits):
-        dtype = jaxext.minimal_unsigned_dtype(splits.size)
+        dtype = minimal_unsigned_dtype(splits.size)
         return jnp.searchsorted(splits, x, **kw).astype(dtype)
 
     return bin_predictors(X, splits)

{bartz-0.5.0.dist-info → bartz-0.7.0.dist-info}/METADATA

@@ -1,14 +1,15 @@
 Metadata-Version: 2.4
 Name: bartz
-Version: 0.5.0
+Version: 0.7.0
 Summary: Super-fast BART (Bayesian Additive Regression Trees) in Python
 Author: Giacomo Petrillo
 Author-email: Giacomo Petrillo <info@giacomopetrillo.com>
 License-Expression: MIT
-Requires-Dist: jax>=0.4.35,<1
-Requires-Dist: jaxlib>=0.4.35,<1
-Requires-Dist: numpy>=1.25.2,<3
-Requires-Dist: scipy>=1.11.4,<2
+Requires-Dist: equinox>=0.12.2
+Requires-Dist: jax>=0.5.3
+Requires-Dist: jaxtyping>=0.3.2
+Requires-Dist: numpy>=1.25.2
+Requires-Dist: scipy>=1.11.4
 Requires-Python: >=3.10
 Project-URL: Documentation, https://gattocrucco.github.io/bartz/docs-dev
 Project-URL: Homepage, https://github.com/Gattocrucco/bartz

bartz-0.7.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+bartz/.DS_Store,sha256=7191af46d7b8c0d4c03c502f94eb01353bc2e615d75c45b3af0e31ab238034b5,6148
+bartz/BART.py,sha256=6b129e20a258d724e0cba5ffaf377b5c4d62e545c1f3a737602c2a0be5d84b96,29601
+bartz/__init__.py,sha256=98c579136a8755390210ada33e713290749e4d7fca58550c791f77f192f4b4a1,1436
+bartz/_version.py,sha256=d3d868979a2f2fa02b20b248259f8c8ac7273a329ae918da1139ad7602695b67,22
+bartz/debug.py,sha256=5082c2dd07f6d3f8353c57491c857cca4aea913647e24502cc375f9aadc84975,43736
+bartz/grove.py,sha256=f64505623feec7edcec96930909f9a9326f290976d7e9bdb7b9abcb32fe425fe,10559
+bartz/jaxext/__init__.py,sha256=6cd2e7c23ccc4f0399fb3f7989312a500fd8bd3f7f07eb85ce537f2c8873f35a,6705
+bartz/jaxext/_autobatch.py,sha256=b5dbaec52e39b4b32c824fde47eaaf33f496a1574cc336d3a79fa71f4c5e348a,7116
+bartz/jaxext/scipy/__init__.py,sha256=a1f5990a75c1c73908565be4cd5fa1c07278ad3a02b78b21e2f1225b388ab6b5,1227
+bartz/jaxext/scipy/special.py,sha256=f0e777c29a77d46d55ff2f0169b3d8581de24d8580b8d6f1c7c1473910fb62d9,8111
+bartz/jaxext/scipy/stats.py,sha256=703beb9fcfe606a195fe9a3143eceed367c87a31bd75763ecdb6e509c2e87f53,1483
+bartz/mcmcloop.py,sha256=e841c076bf5392b1ad8eab955c70a63d18ac321c249d255c7424b62e52c127c1,22134
+bartz/mcmcstep.py,sha256=3b2a4bdad3bbf836efc7cc6f52adac3ba23dce00a1cca4a99e5a8d79d10c3e68,84092
+bartz/prepcovars.py,sha256=50334621ca6ec7a6e35d21ce5ff8b96d4cbf0e696c9bcd10d86268e301a820a7,8728
+bartz-0.7.0.dist-info/WHEEL,sha256=607c46fee47e440c91332c738096ff0f5e54ca3b0818ee85462dd5172a38e793,79
+bartz-0.7.0.dist-info/METADATA,sha256=758e12296acf815c9bae50601ff502e6dee97c9ac0d817232e45e3dfb39665bc,2815
+bartz-0.7.0.dist-info/RECORD,,

{bartz-0.5.0.dist-info → bartz-0.7.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: uv 0.7.4
+Generator: uv 0.7.19
 Root-Is-Purelib: true
 Tag: py3-none-any

bartz/jaxext.py

@@ -1,374 +0,0 @@
-# bartz/src/bartz/jaxext.py
-#
-# Copyright (c) 2024-2025, Giacomo Petrillo
-#
-# This file is part of bartz.
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# 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
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-import functools
-import math
-import warnings
-
-import jax
-from jax import lax, tree_util
-from jax import numpy as jnp
-from scipy import special
-
-
-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
-
-
-def castto(func, type):
-    @functools.wraps(func)
-    def newfunc(*args, **kw):
-        return func(*args, **kw).astype(type)
-
-    return newfunc
-
-
-class scipy:
-    class special:
-        @functools.wraps(special.gammainccinv)
-        def gammainccinv(a, y):
-            a = jnp.asarray(a)
-            y = jnp.asarray(y)
-            shape = jnp.broadcast_shapes(a.shape, y.shape)
-            dtype = float_type(a.dtype, y.dtype)
-            dummy = jax.ShapeDtypeStruct(shape, dtype)
-            ufunc = castto(special.gammainccinv, dtype)
-            return jax.pure_callback(ufunc, dummy, a, y, vmap_method='expand_dims')
-
-    class stats:
-        class invgamma:
-            def ppf(q, a):
-                return 1 / scipy.special.gammainccinv(a, q)
-
-
-def vmap_nodoc(fun, *args, **kw):
-    """
-    Wrapper of `jax.vmap` that preserves the docstring of the input function.
-
-    This is useful if the docstring already takes into account that the
-    arguments have additional axes due to vmap.
-    """
-    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 (excluding
-        unbatched arguments.)
-    in_axes : pytree of int or None, 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. A `None` axis indicates to not batch an argument.
-    out_axes : pytree of ints, default 0
-        The same for outputs (but non-batching is not allowed).
-    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 check_no_nones(axes, tree):
-        def check_not_none(_, axis):
-            assert axis is not None
-
-        tree_util.tree_map(check_not_none, tree, axes)
-
-    def extract_size(axes, tree):
-        def get_size(x, axis):
-            if axis is None:
-                return None
-            else:
-                return x.shape[axis]
-
-        sizes = tree_util.tree_map(get_size, 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 dividend == 0:
-            return min_divisor
-        if min_divisor * min_divisor <= dividend:
-            return next_divisor_small(dividend, min_divisor)
-        return next_divisor_large(dividend, min_divisor)
-
-    def pull_nonbatched(axes, tree):
-        def pull_nonbatched(x, axis):
-            if axis is None:
-                return None
-            else:
-                return x
-
-        return tree_util.tree_map(pull_nonbatched, tree, axes), tree
-
-    def push_nonbatched(axes, tree, original_tree):
-        def push_nonbatched(original_x, x, axis):
-            if axis is None:
-                return original_x
-            else:
-                return x
-
-        return tree_util.tree_map(push_nonbatched, original_tree, tree, axes)
-
-    def move_axes_out(axes, tree):
-        def move_axis_out(x, axis):
-            return jnp.moveaxis(x, axis, 0)
-
-        return tree_util.tree_map(move_axis_out, tree, axes)
-
-    def move_axes_in(axes, tree):
-        def move_axis_in(x, axis):
-            return jnp.moveaxis(x, 0, axis)
-
-        return tree_util.tree_map(move_axis_in, tree, axes)
-
-    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)
-        check_no_nones(out_axes, example_result)
-
-        size = extract_size((in_axes, out_axes), (args, example_result))
-
-        args, nonbatched_args = pull_nonbatched(in_axes, args)
-
-        total_nbytes = sum_nbytes((args, example_result))
-        min_nbatches = total_nbytes // max_io_nbytes + bool(
-            total_nbytes % max_io_nbytes
-        )
-        min_nbatches = max(1, min_nbatches)
-        nbatches = next_divisor(size, min_nbatches)
-        assert 1 <= nbatches <= max(1, 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)
-            args = push_nonbatched(in_axes, args, nonbatched_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__()})'