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

bartz/jaxext.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/jaxext.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -26,11 +26,11 @@ import functools
 import math
 import warnings
 
-from scipy import special
 import jax
+from jax import lax, tree_util
 from jax import numpy as jnp
-from jax import tree_util
-from jax import lax
+from scipy import special
+
 
 def float_type(*args):
     """
@@ -39,62 +39,46 @@ def float_type(*args):
     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
 
-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>`_ """
-    if excluded is None:
-        excluded = ()
-    shape = jnp.broadcast_shapes(*(
-        a.shape
-        for i, a in enumerate(args)
-        if i not in excluded
-    ))
-    ndim = len(shape)
-    padded_args = [
-        a if i in excluded
-        else jnp.expand_dims(a, tuple(range(ndim - a.ndim)))
-        for i, a in enumerate(args)
-    ]
-    result = jax.ShapeDtypeStruct(shape, dtype)
-    return jax.pure_callback(callback, result, *padded_args, vectorized=True, **kwargs)
-
-    # TODO when jax solves this, check version and piggyback on original if new
 
 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 pure_callback_ufunc(ufunc, dtype, a, y)
+            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)
 
-@functools.wraps(jax.vmap)
+
 def vmap_nodoc(fun, *args, **kw):
     """
-    Version of `jax.vmap` that preserves the docstring of the input function.
+    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`.
@@ -114,19 +98,21 @@ def huge_value(x):
     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:
+    if max_value < 2**8:
         return jnp.uint8
-    if max_value < 2 ** 16:
+    if max_value < 2**16:
         return jnp.uint16
-    if max_value < 2 ** 32:
+    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
@@ -144,12 +130,14 @@ def signed_to_unsigned(int_dtype):
     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):
     """
@@ -177,15 +165,18 @@ def unique(x, size, fill_value):
     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.
@@ -222,6 +213,7 @@ def autobatch(func, max_io_nbytes, in_axes=0, out_axes=0, return_nbatches=False)
     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):
@@ -230,6 +222,7 @@ def autobatch(func, max_io_nbytes, in_axes=0, out_axes=0, return_nbatches=False)
                 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)
@@ -238,6 +231,7 @@ def autobatch(func, max_io_nbytes, in_axes=0, out_axes=0, return_nbatches=False)
     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):
@@ -266,6 +260,7 @@ def autobatch(func, max_io_nbytes, in_axes=0, out_axes=0, return_nbatches=False)
                 return None
             else:
                 return x
+
         return tree_util.tree_map(pull_nonbatched, tree, axes), tree
 
     def push_nonbatched(axes, tree, original_tree):
@@ -274,32 +269,38 @@ def autobatch(func, max_io_nbytes, in_axes=0, out_axes=0, return_nbatches=False)
                 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
@@ -319,7 +320,9 @@ def autobatch(func, max_io_nbytes, in_axes=0, out_axes=0, return_nbatches=False)
         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 = 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)
@@ -329,7 +332,9 @@ def autobatch(func, max_io_nbytes, in_axes=0, out_axes=0, return_nbatches=False)
         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}')
+            warnings.warn(
+                f'batch_nbytes = {batch_nbytes} > max_io_nbytes = {max_io_nbytes}'
+            )
 
         def loop(_, args):
             args = move_axes_in(in_axes, args)
@@ -352,10 +357,11 @@ def autobatch(func, max_io_nbytes, in_axes=0, out_axes=0, return_nbatches=False)
 
     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 """
+    """dictionary that acts as a leaf in jax pytrees, to store compile-time
+    values"""
 
     def tree_flatten(self):
         return (), self

bartz/mcmcloop.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/mcmcloop.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -29,22 +29,21 @@ Functions that implement the full BART posterior MCMC loop.
 import functools
 
 import jax
-from jax import random
-from jax import debug
+from jax import debug, lax, random, tree
 from jax import numpy as jnp
-from jax import lax
 
-from . import jaxext
-from . import grove
-from . import mcmcstep
+from . import grove, jaxext, mcmcstep
 
-@functools.partial(jax.jit, static_argnums=(1, 2, 3, 4))
-def run_mcmc(bart, n_burn, n_save, n_skip, callback, key):
+
+@functools.partial(jax.jit, static_argnums=(2, 3, 4, 5))
+def run_mcmc(key, bart, n_burn, n_save, n_skip, callback):
     """
     Run the MCMC for the BART posterior.
 
     Parameters
     ----------
+    key : jax.dtypes.prng_key array
+        The key for random number generation.
     bart : dict
         The initial MCMC state, as created and updated by the functions in
         `bartz.mcmcstep`.
@@ -54,88 +53,123 @@ def run_mcmc(bart, n_burn, n_save, n_skip, callback, key):
         The number of iterations to save.
     n_skip : int
         The number of iterations to skip between each saved iteration, plus 1.
+        The effective burn-in is ``n_burn + n_skip - 1``.
     callback : callable
         An arbitrary function run at each iteration, called with the following
         arguments, passed by keyword:
 
         bart : dict
-            The current MCMC state.
+            The MCMC state just after updating it.
         burnin : bool
             Whether the last iteration was in the burn-in phase.
         i_total : int
             The index of the last iteration (0-based).
         i_skip : int
-            The index of the last iteration, starting from the last saved
-            iteration.
+            The number of MCMC updates from the last saved state. The initial
+            state counts as saved, even if it's not copied into the trace.
         n_burn, n_save, n_skip : int
             The corresponding arguments as-is.
 
         Since this function is called under the jax jit, the values are not
         available at the time the Python code is executed. Use the utilities in
         `jax.debug` to access the values at actual runtime.
-    key : jax.dtypes.prng_key array
-        The key for random number generation.
 
     Returns
     -------
     bart : dict
         The final MCMC state.
-    burnin_trace : dict
+    burnin_trace : dict of (n_burn, ...) arrays
         The trace of the burn-in phase, containing the following subset of
         fields from the `bart` dictionary, with an additional head index that
         runs over MCMC iterations: 'sigma2', 'grow_prop_count',
         'grow_acc_count', 'prune_prop_count', 'prune_acc_count'.
-    main_trace : dict
+    main_trace : dict of (n_save, ...) arrays
         The trace of the main phase, containing the following subset of fields
         from the `bart` dictionary, with an additional head index that runs
         over MCMC iterations: 'leaf_trees', 'var_trees', 'split_trees', plus
         the fields in `burnin_trace`.
+
+    Notes
+    -----
+    The number of MCMC updates is ``n_burn + n_skip * n_save``. The traces do
+    not include the initial state, and include the final state.
     """
 
-    tracelist_burnin = 'sigma2', 'grow_prop_count', 'grow_acc_count', 'prune_prop_count', 'prune_acc_count', 'ratios'
+    tracevars_light = (
+        'sigma2',
+        'grow_prop_count',
+        'grow_acc_count',
+        'prune_prop_count',
+        'prune_acc_count',
+        'ratios',
+    )
+    tracevars_heavy = ('leaf_trees', 'var_trees', 'split_trees')
+
+    def empty_trace(length, bart, tracelist):
+        bart = {k: v for k, v in bart.items() if k in tracelist}
+        return jax.vmap(lambda x: x, in_axes=None, out_axes=0, axis_size=length)(bart)
 
-    tracelist_main = tracelist_burnin + ('leaf_trees', 'var_trees', 'split_trees')
+    trace_light = empty_trace(n_burn + n_save, bart, tracevars_light)
+    trace_heavy = empty_trace(n_save, bart, tracevars_heavy)
 
     callback_kw = dict(n_burn=n_burn, n_save=n_save, n_skip=n_skip)
 
-    def inner_loop(carry, _, tracelist, burnin):
-        bart, i_total, i_skip, key = carry
+    carry = (bart, 0, key, trace_light, trace_heavy)
+
+    def loop(carry, _):
+        bart, i_total, key, trace_light, trace_heavy = carry
+
         key, subkey = random.split(key)
-        bart = mcmcstep.step(bart, subkey)
-        callback(bart=bart, burnin=burnin, i_total=i_total, i_skip=i_skip, **callback_kw)
-        output = {key: bart[key] for key in tracelist if key in bart}
-        return (bart, i_total + 1, i_skip + 1, key), output
-
-    def empty_trace(bart, tracelist):
-        return jax.vmap(lambda x: x, in_axes=None, out_axes=0, axis_size=0)(bart)
-
-    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 = empty_trace(bart, tracelist_burnin)
-
-    def outer_loop(carry, _):
-        bart, i_total, key = carry
-        main_loop = functools.partial(inner_loop, tracelist=[], burnin=False)
-        inner_carry = bart, i_total, 0, key
-        (bart, i_total, _, key), _ = lax.scan(main_loop, inner_carry, None, n_skip)
-        output = {key: bart[key] for key in tracelist_main if key in bart}
-        return (bart, i_total, key), output
-
-    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)
+        bart = mcmcstep.step(subkey, bart)
+
+        burnin = i_total < n_burn
+        i_skip = jnp.where(
+            burnin,
+            i_total + 1,
+            (i_total + 1) % n_skip + jnp.where(i_total + 1 < n_skip, n_burn, 0),
+        )
+        callback(
+            bart=bart, burnin=burnin, i_total=i_total, i_skip=i_skip, **callback_kw
+        )
+
+        i_heavy = jnp.where(burnin, 0, (i_total - n_burn) // n_skip)
+        i_light = jnp.where(burnin, i_total, n_burn + i_heavy)
+
+        def update_trace(index, trace, bart):
+            bart = {k: v for k, v in bart.items() if k in trace}
+
+            def assign_at_index(trace_array, state_array):
+                if trace_array.size:
+                    return trace_array.at[index, ...].set(state_array)
+                else:
+                    # this handles the case where a trace is empty (e.g.,
+                    # no burn-in) because jax refuses to index into an array
+                    # of length 0
+                    return trace_array
+
+            return tree.map(assign_at_index, trace, bart)
+
+        trace_heavy = update_trace(i_heavy, trace_heavy, bart)
+        trace_light = update_trace(i_light, trace_light, bart)
+
+        i_total += 1
+        carry = (bart, i_total, key, trace_light, trace_heavy)
+        return carry, None
+
+    carry, _ = lax.scan(loop, carry, None, n_burn + n_skip * n_save)
+
+    bart, _, _, trace_light, trace_heavy = carry
+
+    burnin_trace = tree.map(lambda x: x[:n_burn, ...], trace_light)
+    main_trace = tree.map(lambda x: x[n_burn:, ...], trace_light)
+    main_trace.update(trace_heavy)
 
     return bart, burnin_trace, main_trace
 
+
 @functools.lru_cache
-    # cache to make the callback function object unique, such that the jit
-    # of run_mcmc recognizes it
+# cache to make the callback function object unique, such that the jit
+# of run_mcmc recognizes it
 def make_simple_print_callback(printevery):
     """
     Create a logging callback function for MCMC iterations.
@@ -150,6 +184,7 @@ def make_simple_print_callback(printevery):
     callback : callable
         A function in the format required by `run_mcmc`.
     """
+
     def callback(*, bart, burnin, i_total, i_skip, n_burn, n_save, n_skip):
         prop_total = len(bart['leaf_trees'])
         grow_prop = bart['grow_prop_count'] / prop_total
@@ -158,18 +193,40 @@ def make_simple_print_callback(printevery):
         prune_acc = bart['prune_acc_count'] / bart['prune_prop_count']
         n_total = n_burn + n_save * n_skip
         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)
+        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(burnin, i_total, n_total, grow_prop, grow_acc, prune_prop, prune_acc, printcond):
+
+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)
-        i_str = str(i_total + 1).rjust(ndigits)
-        print(f'Iteration {i_str}/{total_str} '
+        i_str = str(i_total.item() + 1).rjust(ndigits)
+        # I do i_total.item() + 1 instead of just i_total + 1 to solve a bug
+        # originating when jax is combined with some outdated dependencies. (I
+        # did not track down which dependencies exactly.) Doing .item() makes
+        # the + 1 operation be done by Python instead of by jax. The bug is that
+        # jax hangs completely, with a secondary thread blocked at this line.
+        print(
+            f'Iteration {i_str}/{total_str} '
             f'P_grow={grow_prop:.2f} P_prune={prune_prop:.2f} '
-            f'A_grow={grow_acc:.2f} A_prune={prune_acc:.2f}{burnin_flag}')
+            f'A_grow={grow_acc:.2f} A_prune={prune_acc:.2f}{burnin_flag}'
+        )
+
 
 @jax.jit
 def evaluate_trace(trace, X):
@@ -189,9 +246,13 @@ def evaluate_trace(trace, X):
         The predictions for each iteration of the MCMC.
     """
     evaluate_trees = functools.partial(grove.evaluate_forest, sum_trees=False)
-    evaluate_trees = jaxext.autobatch(evaluate_trees, 2 ** 29, (None, 0, 0, 0))
+    evaluate_trees = jaxext.autobatch(evaluate_trees, 2**29, (None, 0, 0, 0))
+
     def loop(_, state):
-        values = evaluate_trees(X, state['leaf_trees'], state['var_trees'], state['split_trees'])
+        values = evaluate_trees(
+            X, state['leaf_trees'], state['var_trees'], state['split_trees']
+        )
         return None, jnp.sum(values, axis=0, dtype=jnp.float32)
+
     _, y = lax.scan(loop, None, trace)
     return y