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

bartz/jaxext.py

@@ -22,25 +22,25 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+"""Additions to jax."""
+
 import functools
 import math
 import warnings
 
 import jax
-from jax import lax, tree_util
+from jax import lax, random, 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.
-    """
+    """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):
+def _castto(func, type):
     @functools.wraps(func)
     def newfunc(*args, **kw):
         return func(*args, **kw).astype(type)
@@ -49,26 +49,37 @@ def castto(func, type):
 
 
 class scipy:
+    """Mockup of the :external:py:mod:`scipy` module."""
+
     class special:
-        @functools.wraps(special.gammainccinv)
+        """Mockup of the :external:py:mod:`scipy.special` module."""
+
+        @staticmethod
         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 = jax.ShapeDtypeStruct(shape, dtype)
-            ufunc = castto(special.gammainccinv, dtype)
+            ufunc = _castto(special.gammainccinv, dtype)
             return jax.pure_callback(ufunc, dummy, a, y, vmap_method='expand_dims')
 
     class stats:
+        """Mockup of the :external:py:mod:`scipy.stats` module."""
+
         class invgamma:
+            """Class that represents the distribution InvGamma(a, 1)."""
+
+            @staticmethod
             def ppf(q, a):
+                """Percentile point function."""
                 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.
+    Acts like `jax.vmap` but preserves the docstring of the function unchanged.
 
     This is useful if the docstring already takes into account that the
     arguments have additional axes due to vmap.
@@ -99,24 +110,22 @@ def huge_value(x):
         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:
+def minimal_unsigned_dtype(value):
+    """Return the smallest unsigned integer dtype that can represent `value`."""
+    if value < 2**8:
         return jnp.uint8
-    if max_value < 2**16:
+    if value < 2**16:
         return jnp.uint16
-    if max_value < 2**32:
+    if 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.
+    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):
@@ -132,9 +141,7 @@ def signed_to_unsigned(int_dtype):
 
 
 def ensure_unsigned(x):
-    """
-    If x has signed integer type, cast it to the unsigned dtype of the same size.
-    """
+    """If x has signed integer type, cast it to the unsigned dtype of the same size."""
     return x.astype(signed_to_unsigned(x.dtype))
 
 
@@ -358,17 +365,59 @@ 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"""
-
-    def tree_flatten(self):
-        return (), self
+class split:
+    """
+    Split a key into `num` keys.
 
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        return aux_data
+    Parameters
+    ----------
+    key : jax.dtypes.prng_key array
+        The key to split.
+    num : int
+        The number of keys to split into.
+    """
 
-    def __repr__(self):
-        return f'{__class__.__name__}({super().__repr__()})'
+    def __init__(self, key, num=2):
+        self._keys = random.split(key, num)
+
+    def __len__(self):
+        return self._keys.size
+
+    def pop(self, shape=None):
+        """
+        Pop one or more keys from the list.
+
+        Parameters
+        ----------
+        shape : int or tuple of int, optional
+            The shape of the keys to pop. If `None`, a single key is popped.
+            If an integer, that many keys are popped. If a tuple, the keys are
+            reshaped to that shape.
+
+        Returns
+        -------
+        keys : jax.dtypes.prng_key array
+            The popped keys.
+
+        Raises
+        ------
+        IndexError
+            If `shape` is larger than the number of keys left in the list.
+
+        Notes
+        -----
+        The keys are popped from the beginning of the list, so for example
+        ``list(keys.pop(2))`` is equivalent to ``[keys.pop(), keys.pop()]``.
+        """
+        if shape is None:
+            shape = ()
+        elif not isinstance(shape, tuple):
+            shape = (shape,)
+        size_to_pop = math.prod(shape)
+        if size_to_pop > self._keys.size:
+            raise IndexError(
+                f'Cannot pop {size_to_pop} keys from {self._keys.size} keys'
+            )
+        popped_keys = self._keys[:size_to_pop]
+        self._keys = self._keys[size_to_pop:]
+        return popped_keys.reshape(shape)