bartz 0.3.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

bartz/__init__.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,7 +23,7 @@
 # SOFTWARE.
 
 """
-A jax implementation of BART
+Super-fast BART (Bayesian Additive Regression Trees) in Python
 
 See the manual at https://gattocrucco.github.io/bartz/docs
 """

bartz/_version.py

@@ -1 +1 @@
-__version__ = '0.3.0'
+__version__ = '0.4.1'

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.
 #
@@ -45,27 +45,6 @@ def castto(func, type):
         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:
@@ -74,9 +53,11 @@ class scipy:
         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: