Merge pull request #212 from danielward27/mostly_docs

Docs, version bump, and flip Sandwich args

Author: danielward27

flowjax/bijections/__init__.py

@@ -9,6 +9,7 @@
 from .exp import Exp
 from .jax_transforms import Scan, Vmap
 from .masked_autoregressive import MaskedAutoregressive
+from .orthogonal import DiscreteCosine, Householder
 from .planar import Planar
 from .power import Power
 from .rational_quadratic_spline import RationalQuadraticSpline
@@ -26,8 +27,6 @@
     Reshape,
     Sandwich,
 )
-from .utils import EmbedCondition, Flip, Identity, Invert, Permute, Reshape, Sandwich
-from .orthogonal import Householder, DiscreteCosine
 
 __all__ = [
     "AdditiveCondition",

flowjax/bijections/bijection.py

@@ -15,7 +15,6 @@
 from equinox import AbstractVar
 from jaxtyping import Array, ArrayLike
 from paramax import unwrap
-import jax
 
 from flowjax.utils import _get_ufunc_signature, arraylike_to_array
 

flowjax/bijections/chain.py

@@ -2,9 +2,9 @@
 
 from collections.abc import Sequence
 
-from paramax import AbstractUnwrappable, unwrap
 import jax.numpy as jnp
 from jax import Array
+from paramax import AbstractUnwrappable, unwrap
 
 from flowjax.bijections.bijection import AbstractBijection
 from flowjax.utils import check_shapes_match, merge_cond_shapes

flowjax/bijections/coupling.py

@@ -8,7 +8,6 @@
 import equinox as eqx
 import jax.nn as jnn
 import jax.numpy as jnp
-import jax
 import paramax
 from jaxtyping import PRNGKeyArray
 

flowjax/bijections/jax_transforms.py

@@ -3,7 +3,6 @@
 from collections.abc import Callable
 
 import equinox as eqx
-from jax import Array
 import jax.numpy as jnp
 from jax.lax import scan
 from jax.tree_util import tree_leaves, tree_map

flowjax/bijections/orthogonal.py

@@ -1,35 +1,54 @@
-from paramax import AbstractUnwrappable, Parameterize
-from flowjax.bijections.bijection import AbstractBijection
-from jax import Array
+"""Orthogonal transformations."""
+
 import jax.numpy as jnp
-import jax.nn as jnn
+from jax import Array
 from jax.scipy import fft
+from jaxtyping import ArrayLike
+from paramax import AbstractUnwrappable, Parameterize
+
+from flowjax.bijections.bijection import AbstractBijection
+from flowjax.utils import arraylike_to_array
 
 
 class Householder(AbstractBijection):
-    """A Householder reflection bijection.
+    """A Householder reflection.
 
-    This bijection implements a Householder reflection, which is a linear
-    transformation that reflects vectors across a hyperplane defined by a normal
+    A linear transformation reflecting vectors across a hyperplane defined by a normal
     vector (params). The transformation is its own inverse and volume-preserving
-    (determinant = ±1).
+    (determinant = -1). Given a unit vector :math:`v`, the transformation is
+    :math:`y = x - 2(x^T v)v`.
+
+    It is often desirable to stack multiple such transforms (e.g. up to the
+    dimensionality of the data):
 
-    Given a unit vector v, the transformation is:
-    x → x - 2(x·v)v
+    .. doctest::
 
-    Attributes:
-        shape: Shape of the input/output vectors
-        cond_shape: Shape of conditional inputs (None as this bijection is unconditional)
+        >>> from flowjax.bijections import Householder, Scan
+        >>> import jax.random as jr
+        >>> import equinox as eqx
+        >>> import jax.numpy as jnp
+
+        >>> dim = 5
+        >>> keys = jr.split(jr.key(0), dim)
+        >>> householder_stack = Scan(
+        ...    eqx.filter_vmap(lambda key: Householder(jr.normal(key, dim)))(keys)
+        ... )
+
+    Args:
         params: Normal vector defining the reflection hyperplane. The vector is
             normalized in the transformation, so scaling params will have no effect
             on the bijection.
     """
+
     shape: tuple[int, ...]
     unit_vec: Array | AbstractUnwrappable
     cond_shape = None
 
-    def __init__(self, params: Array):
-        self.shape = (params.shape[-1],)
+    def __init__(self, params: ArrayLike):
+        params = arraylike_to_array(params)
+        if params.ndim != 1:
+            raise ValueError("params must be a vector.")
+        self.shape = params.shape
         self.unit_vec = Parameterize(lambda x: x / jnp.linalg.norm(x), params)
 
     def _householder(self, x: Array) -> Array:
@@ -47,11 +66,9 @@ class DiscreteCosine(AbstractBijection):
 
     This bijection applies the DCT or its inverse along a specified axis.
 
-    Attributes:
-        shape: Shape of the input/output arrays
-        cond_shape: Shape of conditional inputs (None as this bijection is unconditional)
-        axis: Axis along which to apply the DCT
-        norm: Normalization method, fixed to 'ortho' to ensure bijectivity
+    Args:
+        shape: Shape of the input/output arrays.
+        axis: Axis along which to apply the DCT.
     """
 
     shape: tuple[int, ...]

flowjax/bijections/softplus.py

@@ -2,15 +2,11 @@
 
 from typing import ClassVar
 
-import jax
 import jax.numpy as jnp
-from jax.nn import softplus, soft_sign
-from jaxtyping import Array, ArrayLike
-from paramax import AbstractUnwrappable, Parameterize, unwrap
-from paramax.utils import inv_softplus
+from jax.nn import softplus
 
 from flowjax.bijections.bijection import AbstractBijection
-from flowjax.utils import arraylike_to_array
+
 
 class SoftPlus(AbstractBijection):
     r"""Transforms to positive domain using softplus :math:`y = \log(1 + \exp(x))`."""

flowjax/bijections/utils.py

@@ -80,10 +80,10 @@ def __init__(self, permutation: Int[Array | np.ndarray, "..."]):
         )
 
     def transform_and_log_det(self, x, condition=None):
-        return x[self.permutation], jnp.array(0.0)
+        return x[self.permutation], jnp.zeros(())
 
     def inverse_and_log_det(self, y, condition=None):
-        return y[self.inverse_permutation], jnp.array(0.0)
+        return y[self.inverse_permutation], jnp.zeros(())
 
 
 class Flip(AbstractBijection):
@@ -308,36 +308,34 @@ def inverse_and_log_det(self, y, condition=None):
 
 
 class Sandwich(AbstractBijection):
-    """A bijection that composes bijections in a nested structure: g⁻¹ ∘ f ∘ g.
-
-    The Sandwich bijection creates a new transformation by "sandwiching" one
-    bijection between the forward and inverse applications of another. Given
-    bijections f and g, it computes:
-        Forward:  x → g⁻¹(f(g(x)))
-        Inverse:  y → g⁻¹(f⁻¹(g(y)))
-
-    This composition pattern is useful for:
-    - Creating symmetries in the transformation
-    - Applying a transformation in a different coordinate system
-    - Building more complex bijections from simpler ones
-
-    Attributes:
-        shape: Shape of the input/output arrays
-        cond_shape: Shape of conditional inputs
-        outer: Transformation g applied first and inverted last
-        inner: Transformation f applied in the middle
+    r"""Composes bijections in a nested structure: :math:`g^{-1} \circ f \circ g`.
+
+    Creates a new transformation by "sandwiching" one bijection between the forward and
+    inverse applications of another. Given bijections :math:`f` and :math:`g`, it
+    computes
+
+    - Forward: :math:`y = g^{-1}(f(g(x)))`
+    - Inverse: :math:`x = g^{-1}(f^{-1}(g(y)))`
+
+    This can be used for e.g. creating symmetries in the transformation or to apply a
+    transformation in a different coordinate system.
+
+    Args:
+        inner: The inner transform.
+        outer: The outer transform.
     """
+
     shape: tuple[int, ...]
     cond_shape: tuple[int, ...] | None
-    outer: AbstractBijection
     inner: AbstractBijection
+    outer: AbstractBijection
 
-    def __init__(self, outer: AbstractBijection, inner: AbstractBijection):
+    def __init__(self, inner: AbstractBijection, outer: AbstractBijection):
         check_shapes_match([outer.shape, inner.shape])
         self.cond_shape = merge_cond_shapes([outer.cond_shape, inner.cond_shape])
         self.shape = inner.shape
-        self.outer = outer
         self.inner = inner
+        self.outer = outer
 
     def transform_and_log_det(self, x: Array, condition=None) -> tuple[Array, Array]:
         chain = Chain([self.outer, self.inner, Invert(self.outer)])

flowjax/flows.py

@@ -32,6 +32,7 @@
     Permute,
     Planar,
     RationalQuadraticSpline,
+    Sandwich,
     Scan,
     TriangularAffine,
     Vmap,
@@ -290,10 +291,9 @@ def triangular_spline_flow(
 ) -> Transformed:
     """Triangular spline flow.
 
-    A single layer consists where each layer consists of a triangular affine
-    transformation with weight normalisation, and an elementwise rational quadratic
-    spline. Tanh is used to constrain to the input to [-1, 1] before spline
-    transformations.
+    Each layer consists of a triangular affine transformation with weight normalisation,
+    and an elementwise rational quadratic spline. Tanh is used to constrain to the input
+    to [-1, 1] before spline transformations.
 
     Args:
         key: Jax random key.
@@ -325,9 +325,7 @@ def make_layer(key):
             lambda t: t.triangular, tri_aff, replace_fn=WeightNormalization
         )
         bijections = [
-            LeakyTanh(tanh_max_val, (dim,)),
-            get_splines(),
-            Invert(LeakyTanh(tanh_max_val, (dim,))),
+            Sandwich(get_splines(), LeakyTanh(tanh_max_val, (dim,))),
             tri_aff,
         ]
 

pyproject.toml

@@ -23,7 +23,7 @@ license = { file = "LICENSE" }
 name = "flowjax"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "17.0.2"
+version = "17.1.0"
 
 [project.urls]
 repository = "https://github.com/danielward27/flowjax"
@@ -34,7 +34,7 @@ dev = [
   "pytest",
   "beartype",
   "ruff",
-  "sphinx",
+  "sphinx <8.2",              # TODO due to https://github.com/tox-dev/sphinx-autodoc-typehints/issues/523
   "sphinx-book-theme",
   "sphinx-copybutton",
   "sphinx-autodoc-typehints",