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

bartz/jaxext/_autobatch.py

@@ -0,0 +1,238 @@
+# bartz/src/bartz/jaxext/_autobatch.py
+#
+# Copyright (c) 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.
+
+"""Implementation of `autobatch`."""
+
+import math
+from collections.abc import Callable
+from functools import wraps
+from warnings import warn
+
+from jax import eval_shape, jit
+from jax import numpy as jnp
+from jax.lax import scan
+from jax.tree import flatten as tree_flatten
+from jax.tree import map as tree_map
+from jax.tree import reduce as tree_reduce
+from jaxtyping import PyTree
+
+
+def expand_axes(axes, tree):
+    """Expand `axes` such that they match the pytreedef of `tree`."""
+
+    def expand_axis(axis, subtree):
+        return tree_map(lambda _: axis, subtree)
+
+    return tree_map(expand_axis, axes, tree, is_leaf=lambda x: x is None)
+
+
+def check_no_nones(axes, tree):
+    def check_not_none(_, axis):
+        assert axis is not None
+
+    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_map(get_size, tree, axes)
+    sizes, _ = 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_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_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_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_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_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_map(batch, tree)
+
+
+def unbatch(tree):
+    def unbatch(x):
+        return x.reshape((x.shape[0] * x.shape[1],) + x.shape[2:])
+
+    return 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_map(check_same, tree1, tree2)
+
+
+def autobatch(
+    func: Callable,
+    max_io_nbytes: int,
+    in_axes: PyTree[int | None] = 0,
+    out_axes: PyTree[int] = 0,
+    return_nbatches: bool = False,
+) -> Callable:
+    """
+    Batch a function such that each batch is smaller than a threshold.
+
+    Parameters
+    ----------
+    func
+        A jittable function with positional arguments only, with inputs and
+        outputs pytrees of arrays.
+    max_io_nbytes
+        The maximum number of input + output bytes in each batch (excluding
+        unbatched arguments.)
+    in_axes
+        A tree matching (a prefix of) the structure of the function input,
+        indicating along which axes each array should be batched. A `None` axis
+        indicates to not batch an argument.
+    out_axes
+        The same for outputs (but non-batching is not allowed).
+    return_nbatches
+        If True, the number of batches is returned as a second output.
+
+    Returns
+    -------
+    A function with the same signature as `func`, save for the return value if `return_nbatches`.
+    """
+    initial_in_axes = in_axes
+    initial_out_axes = out_axes
+
+    @jit
+    @wraps(func)
+    def batched_func(*args):
+        example_result = 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
+            msg = f'batch_nbytes = {batch_nbytes} > max_io_nbytes = {max_io_nbytes}'
+            warn(msg)
+
+        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 = 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

bartz/jaxext/scipy/__init__.py

@@ -0,0 +1,25 @@
+# bartz/src/bartz/jaxext/scipy/__init__.py
+#
+# Copyright (c) 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.
+
+"""Mockup of the :external:py:mod:`scipy` module."""

bartz/jaxext/scipy/special.py

@@ -0,0 +1,240 @@
+# bartz/src/bartz/jaxext/scipy/special.py
+#
+# Copyright (c) 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.
+
+"""Mockup of the :external:py:mod:`scipy.special` module."""
+
+from functools import wraps
+
+from jax import ShapeDtypeStruct, pure_callback
+from jax import numpy as jnp
+from scipy.special import gammainccinv as scipy_gammainccinv
+
+
+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, dtype):
+    @wraps(func)
+    def newfunc(*args, **kw):
+        return func(*args, **kw).astype(dtype)
+
+    return newfunc
+
+
+def gammainccinv(a, y):
+    """Survival function inverse of the Gamma(a, 1) distribution."""
+    a = jnp.asarray(a)
+    y = jnp.asarray(y)
+    shape = jnp.broadcast_shapes(a.shape, y.shape)
+    dtype = _float_type(a.dtype, y.dtype)
+    dummy = ShapeDtypeStruct(shape, dtype)
+    ufunc = _castto(scipy_gammainccinv, dtype)
+    return pure_callback(ufunc, dummy, a, y, vmap_method='expand_dims')
+
+
+################# COPIED AND ADAPTED FROM JAX ##################
+# Copyright 2018 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+from jax import debug_infs, lax
+
+
+def ndtri(p):
+    """Compute the inverse of the CDF of the Normal distribution function.
+
+    This is a patch of `jax.scipy.special.ndtri`.
+    """
+    dtype = lax.dtype(p)
+    if dtype not in (jnp.float32, jnp.float64):
+        msg = f'x.dtype={dtype} is not supported, see docstring for supported types.'
+        raise TypeError(msg)
+    return _ndtri(p)
+
+
+def _ndtri(p):
+    # Constants used in piece-wise rational approximations. Taken from the cephes
+    # library:
+    # https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
+    p0 = list(
+        reversed(
+            [
+                -5.99633501014107895267e1,
+                9.80010754185999661536e1,
+                -5.66762857469070293439e1,
+                1.39312609387279679503e1,
+                -1.23916583867381258016e0,
+            ]
+        )
+    )
+    q0 = list(
+        reversed(
+            [
+                1.0,
+                1.95448858338141759834e0,
+                4.67627912898881538453e0,
+                8.63602421390890590575e1,
+                -2.25462687854119370527e2,
+                2.00260212380060660359e2,
+                -8.20372256168333339912e1,
+                1.59056225126211695515e1,
+                -1.18331621121330003142e0,
+            ]
+        )
+    )
+    p1 = list(
+        reversed(
+            [
+                4.05544892305962419923e0,
+                3.15251094599893866154e1,
+                5.71628192246421288162e1,
+                4.40805073893200834700e1,
+                1.46849561928858024014e1,
+                2.18663306850790267539e0,
+                -1.40256079171354495875e-1,
+                -3.50424626827848203418e-2,
+                -8.57456785154685413611e-4,
+            ]
+        )
+    )
+    q1 = list(
+        reversed(
+            [
+                1.0,
+                1.57799883256466749731e1,
+                4.53907635128879210584e1,
+                4.13172038254672030440e1,
+                1.50425385692907503408e1,
+                2.50464946208309415979e0,
+                -1.42182922854787788574e-1,
+                -3.80806407691578277194e-2,
+                -9.33259480895457427372e-4,
+            ]
+        )
+    )
+    p2 = list(
+        reversed(
+            [
+                3.23774891776946035970e0,
+                6.91522889068984211695e0,
+                3.93881025292474443415e0,
+                1.33303460815807542389e0,
+                2.01485389549179081538e-1,
+                1.23716634817820021358e-2,
+                3.01581553508235416007e-4,
+                2.65806974686737550832e-6,
+                6.23974539184983293730e-9,
+            ]
+        )
+    )
+    q2 = list(
+        reversed(
+            [
+                1.0,
+                6.02427039364742014255e0,
+                3.67983563856160859403e0,
+                1.37702099489081330271e0,
+                2.16236993594496635890e-1,
+                1.34204006088543189037e-2,
+                3.28014464682127739104e-4,
+                2.89247864745380683936e-6,
+                6.79019408009981274425e-9,
+            ]
+        )
+    )
+
+    dtype = lax.dtype(p).type
+    shape = jnp.shape(p)
+
+    def _create_polynomial(var, coeffs):
+        """Compute n_th order polynomial via Horner's method."""
+        coeffs = np.array(coeffs, dtype)
+        if not coeffs.size:
+            return jnp.zeros_like(var)
+        return coeffs[0] + _create_polynomial(var, coeffs[1:]) * var
+
+    maybe_complement_p = jnp.where(p > dtype(-np.expm1(-2.0)), dtype(1.0) - p, p)
+    # Write in an arbitrary value in place of 0 for p since 0 will cause NaNs
+    # later on. The result from the computation when p == 0 is not used so any
+    # number that doesn't result in NaNs is fine.
+    sanitized_mcp = jnp.where(
+        maybe_complement_p == dtype(0.0),
+        jnp.full(shape, dtype(0.5)),
+        maybe_complement_p,
+    )
+
+    # Compute x for p > exp(-2): x/sqrt(2pi) = w + w**3 P0(w**2)/Q0(w**2).
+    w = sanitized_mcp - dtype(0.5)
+    ww = lax.square(w)
+    x_for_big_p = w + w * ww * (_create_polynomial(ww, p0) / _create_polynomial(ww, q0))
+    x_for_big_p *= -dtype(np.sqrt(2.0 * np.pi))
+
+    # Compute x for p <= exp(-2): x = z - log(z)/z - (1/z) P(1/z) / Q(1/z),
+    # where z = sqrt(-2. * log(p)), and P/Q are chosen between two different
+    # arrays based on whether p < exp(-32).
+    z = lax.sqrt(dtype(-2.0) * lax.log(sanitized_mcp))
+    first_term = z - lax.log(z) / z
+    second_term_small_p = (
+        _create_polynomial(dtype(1.0) / z, p2)
+        / _create_polynomial(dtype(1.0) / z, q2)
+        / z
+    )
+    second_term_otherwise = (
+        _create_polynomial(dtype(1.0) / z, p1)
+        / _create_polynomial(dtype(1.0) / z, q1)
+        / z
+    )
+    x_for_small_p = first_term - second_term_small_p
+    x_otherwise = first_term - second_term_otherwise
+
+    x = jnp.where(
+        sanitized_mcp > dtype(np.exp(-2.0)),
+        x_for_big_p,
+        jnp.where(z >= dtype(8.0), x_for_small_p, x_otherwise),
+    )
+
+    x = jnp.where(p > dtype(1.0 - np.exp(-2.0)), x, -x)
+    with debug_infs(False):
+        infinity = jnp.full(shape, dtype(np.inf))
+        neg_infinity = -infinity
+    return jnp.where(
+        p == dtype(0.0), neg_infinity, jnp.where(p == dtype(1.0), infinity, x)
+    )
+
+
+################################################################

bartz/jaxext/scipy/stats.py

@@ -0,0 +1,36 @@
+# bartz/src/bartz/jaxext/scipy/stats.py
+#
+# Copyright (c) 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.
+
+"""Mockup of the :external:py:mod:`scipy.stats` module."""
+
+from bartz.jaxext.scipy.special import gammainccinv
+
+
+class invgamma:
+    """Class that represents the distribution InvGamma(a, 1)."""
+
+    @staticmethod
+    def ppf(q, a):
+        """Percentile point function."""
+        return 1 / gammainccinv(a, q)