bartz 0.0.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

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:
@@ -100,15 +100,24 @@ def run_mcmc(bart, n_burn, n_save, n_skip, callback, key):
     def inner_loop(carry, _, tracelist, burnin):
         bart, i_total, i_skip, key = carry
         key, subkey = random.split(key)
-        bart = mcmcstep.mcmc_step(bart, subkey)
+        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}
         return (bart, i_total + 1, i_skip + 1, key), output
 
-    # TODO avoid invoking this altogether if burnin is 0 to shorten compilation time & size
-    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)
+    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 = empty_trace(bart, tracelist_burnin)
 
     def outer_loop(carry, _):
         bart, i_total, key = carry
@@ -118,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
 
@@ -127,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.
@@ -149,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)
@@ -180,6 +194,6 @@ def evaluate_trace(trace, X):
         The predictions for each iteration of the MCMC.
     """
     def loop(_, state):
-        return None, grove.evaluate_tree_vmap_x(X, state['leaf_trees'], state['var_trees'], state['split_trees'], jnp.float32)
+        return None, grove.evaluate_forest(X, state['leaf_trees'], state['var_trees'], state['split_trees'], jnp.float32)
     _, y = lax.scan(loop, None, trace)
     return y