Add stochastic base models and flatten model

Update MLPModel to MLP

Author: RunnersNum40

oryx/models/__init__.py

@@ -4,15 +4,21 @@
 Models take inputs and produce outputs, and may have state.
 """
 
-from oryx.models.base_model import AbstractModel, AbstractStatefulModel
-from oryx.models.mlp import MLPModel
-from oryx.models.ncde import (
+from .base_model import (
+    AbstractModel,
+    AbstractStatefulModel,
+    AbstractStochasticStatefulModel,
+    AbstractStochaticModel,
+)
+from .flatten import Flatten
+from .mlp import MLP
+from .ncde import (
     AbstractNCDETerm,
     AbstractNeuralCDE,
     MLPNCDETerm,
     MLPNeuralCDE,
 )
-from oryx.models.node import (
+from .node import (
     AbstractNeuralODE,
     AbstractNODETerm,
     MLPNeuralODE,
@@ -22,6 +28,9 @@
 __all__ = [
     "AbstractModel",
     "AbstractStatefulModel",
+    "AbstractStochasticStatefulModel",
+    "AbstractStochaticModel",
+    "Flatten",
     "AbstractNeuralODE",
     "AbstractNODETerm",
     "MLPNeuralODE",
@@ -30,5 +39,5 @@
     "AbstractNCDETerm",
     "MLPNeuralCDE",
     "MLPNCDETerm",
-    "MLPModel",
+    "MLP",
 ]

oryx/models/base_model.py

@@ -2,6 +2,7 @@
 from typing import Concatenate
 
 import equinox as eqx
+from jaxtyping import Key
 
 
 class AbstractModel[**InType, OutType](eqx.Module, strict=True):
@@ -25,3 +26,41 @@ def __call__(
         self, state: eqx.nn.State, *args: InType.args, **kwargs: InType.kwargs
     ) -> tuple[eqx.nn.State, *OutType]:
         """Return an output given inputs and the state."""
+
+
+class AbstractStochaticModel[**InType, OutType](
+    AbstractModel[
+        Concatenate[Key, InType],
+        OutType,
+    ],
+    strict=True,
+):
+    """Base class for stochastic models that take inputs and produce outputs."""
+
+    # TODO: Revisit this if Python ever adds support for concatenating keyword
+    # parameters, I'd rather key as a keyword to fit the style of the library
+    @abstractmethod
+    def __call__(
+        self, key: Key, *args: InType.args, **kwargs: InType.kwargs
+    ) -> OutType:
+        """Return an output given an input."""
+
+
+class AbstractStochasticStatefulModel[**InType, *OutType](
+    AbstractStatefulModel[
+        Concatenate[Key, InType],
+        *OutType,
+    ],
+    AbstractStochaticModel[
+        Concatenate[eqx.nn.State, InType],
+        tuple[eqx.nn.State, *OutType],
+    ],
+    strict=True,
+):
+    """Base class for stochastic models with state."""
+
+    @abstractmethod
+    def __call__(
+        self, state: eqx.nn.State, key: Key, *args: InType.args, **kwargs: InType.kwargs
+    ) -> tuple[eqx.nn.State, *OutType]:
+        """Return an output given inputs and the state."""

oryx/models/flatten.py

@@ -0,0 +1,14 @@
+from jax import numpy as jnp
+from jaxtyping import Array, Float
+
+from .base_model import AbstractModel
+
+
+class Flatten(
+    AbstractModel[[Float[Array, " ..."]], Float[Array, " out_size"]], strict=True
+):
+    """Warps the JAX `jnp.ravel` function to flatten an input array."""
+
+    def __call__(self, x: Float[Array, " ..."]) -> Float[Array, " out_size"]:
+        """Flatten the input array into a 1D array."""
+        return jnp.ravel(x)

oryx/models/mlp.py

@@ -7,7 +7,7 @@
 from .base_model import AbstractModel
 
 
-class MLPModel(
+class MLP(
     AbstractModel[[Float[Array, " in_size"]], Float[Array, " out_size"]], strict=True
 ):
     """Wrapper around eqx.nn.MLP."""

tests/test_models.py

@@ -4,8 +4,14 @@
 from jax import numpy as jnp
 from jax import random as jr
 
-from oryx.models.base_model import AbstractModel, AbstractStatefulModel
-from oryx.models.mlp import MLPModel
+from oryx.models.base_model import (
+    AbstractModel,
+    AbstractStatefulModel,
+    AbstractStochasticStatefulModel,
+    AbstractStochaticModel,
+)
+from oryx.models.flatten import Flatten
+from oryx.models.mlp import MLP
 from oryx.models.ncde.ncde import MLPNeuralCDE
 from oryx.models.ncde.term import AbstractNCDETerm, MLPNCDETerm
 from oryx.models.node.node import MLPNeuralODE
@@ -54,7 +60,7 @@ def __call__(self, state):
 def test_mlpmodel_forward_shape_and_determinism():
     key = jr.key(0)
     in_size, out_size = 7, 4
-    m = MLPModel(in_size=in_size, out_size=out_size, width_size=10, depth=2, key=key)
+    m = MLP(in_size=in_size, out_size=out_size, width_size=10, depth=2, key=key)
     x = jnp.zeros(in_size)
     y1 = m(x)
     y2 = m(x)
@@ -65,14 +71,14 @@ def test_mlpmodel_forward_shape_and_determinism():
 
 def test_mlpmodel_wrong_input_shape_raises():
     key = jr.key(1)
-    m = MLPModel(in_size=5, out_size=3, width_size=8, depth=1, key=key)
+    m = MLP(in_size=5, out_size=3, width_size=8, depth=1, key=key)
     x_bad = jnp.ones((6,))
     with pytest.raises(Exception):
         _ = m(x_bad)
 
 
 @pytest.mark.parametrize("time_in_input", [True, False])
-def test_mlpneuralode_solve_and_call(time_in_input):
+def test_mlpneuralode_solve_and_call(time_in_input: bool):
     key = jr.key(4)
     in_size, out_size, latent_size = 3, 2, 4
     model = MLPNeuralODE(
@@ -138,7 +144,7 @@ class NoCallNCDETerm(AbstractNCDETerm):
 
 
 @pytest.mark.parametrize("add_time", [True, False])
-def test_mlpncdeterm_output_shape(add_time):
+def test_mlpncdeterm_output_shape(add_time: bool):
     key = jr.key(6)
     input_size, data_size = 3, 5
     term = MLPNCDETerm(
@@ -157,7 +163,7 @@ def test_mlpncdeterm_output_shape(add_time):
 
 
 @pytest.mark.parametrize("inference", [True, False])
-def test_mlpneuralcde(inference):
+def test_mlpneuralcde(inference: bool):
     key = jr.key(7)
     in_size, out_size, latent_size, state_size = 4, 2, 6, 4
     model, state = eqx.nn.make_with_state(MLPNeuralCDE)(
@@ -191,7 +197,7 @@ def test_mlpneuralcde(inference):
 
 
 @pytest.mark.parametrize("time_in_input", [True, False])
-def test_ncde_z0_and_coeffs_shapes(time_in_input):
+def test_ncde_z0_and_coeffs_shapes(time_in_input: bool):
     key = jr.key(0)
     in_size, latent_size = 3, 5
     model = MLPNeuralCDE(
@@ -219,7 +225,7 @@ def test_ncde_z0_and_coeffs_shapes(time_in_input):
 
 
 @pytest.mark.parametrize("inference", [True, False])
-def test_ncde_solve_output_shape(inference):
+def test_ncde_solve_output_shape(inference: bool):
     key = jr.key(1)
     in_size, latent_size = 2, 4
     model = MLPNeuralCDE(
@@ -242,7 +248,7 @@ def test_ncde_solve_output_shape(inference):
 
 
 @pytest.mark.parametrize("inference", [True, False])
-def test_ncde_state_rollover(inference):
+def test_ncde_state_rollover(inference: bool):
     key = jr.key(2)
     in_size, out_size, latent_size = 2, 1, 3
     model, state = eqx.nn.make_with_state(MLPNeuralCDE)(
@@ -266,3 +272,17 @@ def test_ncde_state_rollover(inference):
     assert not jnp.isnan(ts).any()
     assert jnp.allclose(ts, jnp.array([1.0, 2.0, 3.0]))
     assert model.t1(state) == 3.0
+
+
+@pytest.mark.parametrize("dims", range(4))
+def test_flatten_model(dims: int):
+    model = Flatten()
+
+    input_shape = (2,) * dims
+    input_data = jnp.ones(input_shape)
+
+    output = model(input_data)
+
+    assert isinstance(output, jnp.ndarray)
+    assert output.shape == (2**dims,)
+    assert jnp.array_equal(output, jnp.ones((2**dims,)))