Merge pull request #158 from danielward27/types

Types

flowjax/bijections/affine.py

@@ -67,7 +67,7 @@ class Loc(AbstractBijection):
     cond_shape: ClassVar[None] = None
 
     def __init__(self, loc: ArrayLike):
-        self.loc = arraylike_to_array(loc)
+        self.loc = arraylike_to_array(loc, dtype=float)
         self.shape = self.loc.shape
 
     def transform(self, x, condition=None):
@@ -98,6 +98,7 @@ def __init__(
         self,
         scale: ArrayLike,
     ):
+        scale = arraylike_to_array(scale, "scale", dtype=float)
         self.scale = wrappers.BijectionReparam(scale, SoftPlus())
         self.shape = jnp.shape(wrappers.unwrap(scale))
 

flowjax/bijections/utils.py

@@ -6,7 +6,8 @@
 
 import equinox as eqx
 import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike
+import numpy as np
+from jaxtyping import Array, Int
 
 from flowjax.bijections.bijection import AbstractBijection
 from flowjax.utils import arraylike_to_array
@@ -63,11 +64,11 @@ class Permute(AbstractBijection):
     permutation: tuple[Array, ...]
     inverse_permutation: tuple[Array, ...]
 
-    def __init__(self, permutation: ArrayLike):
-        permutation = arraylike_to_array(permutation)
+    def __init__(self, permutation: Int[Array | np.ndarray, "..."]):
+        permutation = arraylike_to_array(permutation, dtype=int)
         permutation = eqx.error_if(
             permutation,
-            permutation.ravel().sort() != jnp.arange(permutation.size),
+            permutation.ravel().sort() != jnp.arange(permutation.size, dtype=int),
             "Invalid permutation array provided.",
         )
         self.shape = permutation.shape

flowjax/distributions.py

@@ -500,8 +500,7 @@ class Uniform(AbstractLocScaleDistribution):
     bijection: Affine
 
     def __init__(self, minval: ArrayLike, maxval: ArrayLike):
-        minval, maxval = arraylike_to_array(minval), arraylike_to_array(maxval)
-        shape = jnp.broadcast_shapes(minval.shape, maxval.shape)
+        shape = jnp.broadcast_shapes(jnp.shape(minval), jnp.shape(maxval))
         minval, maxval = eqx.error_if(
             (minval, maxval), maxval <= minval, "minval must be less than the maxval."
         )
@@ -596,7 +595,7 @@ class _StandardStudentT(AbstractDistribution):
     df: Array | AbstractUnwrappable[Array]
 
     def __init__(self, df: ArrayLike):
-        df = arraylike_to_array(df)
+        df = arraylike_to_array(df, dtype=float)
         df = eqx.error_if(df, df <= 0, "Degrees of freedom values must be positive.")
         self.shape = jnp.shape(df)
         self.df = BijectionReparam(df, SoftPlus())

flowjax/experimental/numpyro.py

@@ -209,7 +209,10 @@ def condition(self):
         return jax.lax.stop_gradient(self._condition)
 
     def tree_flatten(self):
-        raise NotImplementedError()
+        return (self.bijection, self._condition, self.domain, self.codomain), (
+            ("bijection", "_condition", "domain", "codomain"),
+            {},
+        )
 
     def _argcheck_domains(self):
         for k, v in {"domain": self.domain, "codomain": self.codomain}.items():