Remove equinox linear workaround

Author: aseyboldt

python/nutpie/normalizing_flow.py

@@ -10,6 +10,7 @@
 import flowjax.flows
 import numpy as np
 from paramax import Parameterize
+from equinox.nn import Linear
 
 
 def _generate_sequences(k, r_vals):
@@ -115,102 +116,6 @@ def _generate_permutations(rng, n_dim, n_layers, max_run=3):
     return permutations.T, is_in_first.sum(0)
 
 
-# Fix upstream bug for zero-dimensional arrays
-class Linear(eqx.Module, strict=True):
-    """Performs a linear transformation."""
-
-    weight: jax.Array
-    bias: jax.Array | None
-    in_features: Union[int, Literal["scalar"]] = eqx.field(static=True)
-    out_features: Union[int, Literal["scalar"]] = eqx.field(static=True)
-    use_bias: bool = eqx.field(static=True)
-
-    def __init__(
-        self,
-        in_features: Union[int, Literal["scalar"]],
-        out_features: Union[int, Literal["scalar"]],
-        use_bias: bool = True,
-        dtype=None,
-        *,
-        key,
-    ):
-        """**Arguments:**
-
-        - `in_features`: The input size. The input to the layer should be a vector of
-            shape `(in_features,)`
-        - `out_features`: The output size. The output from the layer will be a vector
-            of shape `(out_features,)`.
-        - `use_bias`: Whether to add on a bias as well.
-        - `dtype`: The dtype to use for the weight and the bias in this layer.
-            Defaults to either `jax.numpy.float32` or `jax.numpy.float64` depending
-            on whether JAX is in 64-bit mode.
-        - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
-            initialisation. (Keyword only argument.)
-
-        Note that `in_features` also supports the string `"scalar"` as a special value.
-        In this case the input to the layer should be of shape `()`.
-
-        Likewise `out_features` can also be a string `"scalar"`, in which case the
-        output from the layer will have shape `()`.
-        """
-        dtype = np.float32 if dtype is None else dtype
-        wkey, bkey = jax.random.split(key, 2)
-        in_features_ = 1 if in_features == "scalar" else in_features
-        out_features_ = 1 if out_features == "scalar" else out_features
-        if in_features_ == 0:
-            lim = 1.0
-        else:
-            lim = 1 / math.sqrt(in_features_)
-        wshape = (out_features_, in_features_)
-        self.weight = eqx.nn._misc.default_init(wkey, wshape, dtype, lim)
-        bshape = (out_features_,)
-        self.bias = (
-            eqx.nn._misc.default_init(bkey, bshape, dtype, lim) if use_bias else None
-        )
-
-        self.in_features = in_features
-        self.out_features = out_features
-        self.use_bias = use_bias
-
-    @jax.named_scope("eqx.nn.Linear")
-    def __call__(self, x: jax.Array, *, key=None) -> jax.Array:
-        """**Arguments:**
-
-        - `x`: The input. Should be a JAX array of shape `(in_features,)`. (Or shape
-            `()` if `in_features="scalar"`.)
-        - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
-            (Keyword only argument.)
-
-        !!! info
-
-            If you want to use higher order tensors as inputs (for example featuring "
-            "batch dimensions) then use `jax.vmap`. For example, for an input `x` of "
-            "shape `(batch, in_features)`, using
-            ```python
-            linear = equinox.nn.Linear(...)
-            jax.vmap(linear)(x)
-            ```
-            will produce the appropriate output of shape `(batch, out_features)`.
-
-        **Returns:**
-
-        A JAX array of shape `(out_features,)`. (Or shape `()` if
-        `out_features="scalar"`.)
-        """
-
-        if self.in_features == "scalar":
-            if jnp.shape(x) != ():
-                raise ValueError("x must have scalar shape")
-            x = jnp.broadcast_to(x, (1,))
-        x = self.weight @ x
-        if self.bias is not None:
-            x = x + self.bias
-        if self.out_features == "scalar":
-            assert jnp.shape(x) == (1,)
-            x = jnp.squeeze(x)
-        return x
-
-
 class FactoredMLP(eqx.Module, strict=True):
     """Standard Multi-Layer Perceptron; also known as a feed-forward network.