diff --git a/.github/workflows/python_app.yaml b/.github/workflows/python_app.yaml index 9b4dab0..44fe1fa 100644 --- a/.github/workflows/python_app.yaml +++ b/.github/workflows/python_app.yaml @@ -19,14 +19,14 @@ jobs: steps: - uses: actions/checkout@v3 - - name: Set up Python 3.13 + - name: Set up Python 3.14 uses: actions/setup-python@v3 with: - python-version: "3.13" + python-version: "3.14" - name: Install dependencies run: | python -m pip install --upgrade pip - pip install pytest pytest-cov + pip install pytest pip install -e ".[all]" - name: Test with pytest run: | diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8adf55d..f24593f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,15 +1,15 @@ repos: - repo: https://github.com/psf/black - rev: 23.3.0 + rev: 26.5.1 hooks: - id: black args: [--line-length=88] - repo: https://github.com/pycqa/isort - rev: 5.13.2 + rev: 8.0.1 hooks: - id: isort - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.3.0 + rev: v6.0.0 hooks: - id: trailing-whitespace - id: end-of-file-fixer diff --git a/memax/equinox/set_actions/lstm.py b/memax/equinox/set_actions/lstm.py index 6930fb4..e068b3f 100644 --- a/memax/equinox/set_actions/lstm.py +++ b/memax/equinox/set_actions/lstm.py @@ -1,21 +1,21 @@ -from beartype.typing import Callable, Optional, Tuple - +import equinox as eqx import jax import jax.numpy as jnp from beartype import beartype as typechecker +from beartype.typing import Callable, Optional, Tuple from equinox import nn from jaxtyping import Array, Float, PRNGKeyArray, Shaped, jaxtyped -from memax.equinox.groups import BinaryAlgebra, SetAction, Resettable from memax.equinox.gras import GRAS -from memax.mtypes import Input, InputEmbedding, StartFlag +from memax.equinox.groups import BinaryAlgebra, Resettable, SetAction from memax.equinox.scans import set_action_scan +from memax.mtypes import Input, InputEmbedding, StartFlag LSTMRecurrentState = Tuple[Float[Array, "Recurrent"], Float[Array, "Recurrent"]] LSTMRecurrentStateWithReset = Tuple[LSTMRecurrentState, StartFlag] -class LSTMMagma(SetAction): +class LSTMSetAction(SetAction): """ The Long Short-Term Memory set action @@ -32,7 +32,11 @@ class LSTMMagma(SetAction): W_o: nn.Linear W_c: nn.Linear - def __init__(self, recurrent_size: int, key): + def __init__( + self, + recurrent_size: int, + key, + ): self.recurrent_size = recurrent_size keys = jax.random.split(key, 8) self.U_f = nn.Linear( @@ -65,7 +69,7 @@ def __call__( f_c = jax.nn.tanh(self.W_c(x_t) + self.U_c(h)) c = f_f * c + f_i * f_c - h = f_o * c + h = f_o * jax.nn.tanh(c) return (c, h) @@ -101,7 +105,7 @@ class LSTM(GRAS): def __init__(self, recurrent_size, key): keys = jax.random.split(key, 3) - self.algebra = Resettable(LSTMMagma(recurrent_size, key=keys[0])) + self.algebra = Resettable(LSTMSetAction(recurrent_size, key=keys[0])) self.scan = set_action_scan @jaxtyped(typechecker=typechecker) @@ -118,7 +122,7 @@ def backward_map( h: LSTMRecurrentStateWithReset, x: Input, key: Optional[Shaped[PRNGKeyArray, ""]] = None, - ) -> Float[Array, "Recurrent"]: + ) -> Float[Array, "Recurrent"]: (c_t, h_t), reset_flag = h emb, start = x return h_t diff --git a/setup.py b/setup.py index 1b042a0..8b4af54 100644 --- a/setup.py +++ b/setup.py @@ -18,31 +18,29 @@ "beartype", ], extras_require={ - 'equinox': ['equinox'], + "equinox": ["equinox"], # TODO: Update if flax fixes their shit - 'flax': [ - 'flax', - 'please-downgrade-to-python-3.13-for-flax; python_version >= "3.14"', + "flax": [ + "flax", ], - 'train': [ - 'datasets', - 'tqdm', - 'pillow', - 'wandb', + "train": [ + "datasets", + "tqdm", + "pillow", + "wandb", ], - 'all': [ - 'equinox', - 'flax', - 'please-downgrade-to-python-3.13-for-flax; python_version >= "3.14"', + "all": [ + "equinox", + "flax", # train - 'datasets', - 'tqdm', - 'pillow', - 'wandb', + "datasets", + "tqdm", + "pillow", + "wandb", + ], + "test": [ + "pytest", ], - 'test': [ - 'pytest', - ] }, classifiers=[ "Programming Language :: Python :: 3", diff --git a/tests/test_continuous_localization.py b/tests/test_continuous_localization.py index 969c7e6..65fbf79 100644 --- a/tests/test_continuous_localization.py +++ b/tests/test_continuous_localization.py @@ -1,35 +1,53 @@ -from memax.datasets.continuous_localization import step +import jax import jax.numpy as jnp from jax.scipy.spatial.transform import Rotation -import jax + +from memax.datasets.continuous_localization import step + def test_step(): dx = jnp.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]]) x = jnp.array([[1, 0, 0], [2, 0, 0], [2, 1, 0]]) - drot = Rotation.from_quat(jnp.stack([ - Rotation.from_euler('z', jnp.pi/2).as_quat(), - Rotation.from_euler('y', -jnp.pi/2).as_quat(), - Rotation.from_euler('YZ', jnp.array([jnp.pi/2, -jnp.pi/2])).as_quat() - ], axis=0)) - rot = Rotation.from_quat(jnp.stack([ - Rotation.from_euler('z', jnp.pi / 2).as_quat(), - Rotation.from_euler('zx', jnp.array([jnp.pi/2, jnp.pi/2])).as_quat(), - Rotation.identity().as_quat(), - ], axis=0)) - - x_start = jnp.zeros((1,3)) + drot = Rotation.from_quat( + jnp.stack( + [ + Rotation.from_euler("z", jnp.pi / 2).as_quat(), + Rotation.from_euler("y", -jnp.pi / 2).as_quat(), + Rotation.from_euler( + "YZ", jnp.array([jnp.pi / 2, -jnp.pi / 2]) + ).as_quat(), + ], + axis=0, + ) + ) + rot = Rotation.from_quat( + jnp.stack( + [ + Rotation.from_euler("z", jnp.pi / 2).as_quat(), + Rotation.from_euler( + "zx", jnp.array([jnp.pi / 2, jnp.pi / 2]) + ).as_quat(), + Rotation.identity().as_quat(), + ], + axis=0, + ) + ) + + x_start = jnp.zeros((1, 3)) rot_start = jax.vmap(Rotation.identity, axis_size=1)() - _, (x_pred, rot_pred) = jax.lax.scan(step, (x_start, rot_start), (dx[:,None], drot[:,None])) + _, (x_pred, rot_pred) = jax.lax.scan( + step, (x_start, rot_start), (dx[:, None], drot[:, None]) + ) - pred_rot = rot_pred.as_matrix()[:,0] + pred_rot = rot_pred.as_matrix()[:, 0] true_rot = rot.as_matrix() - pred_x = x_pred[:,0] - true_x = x + pred_x = x_pred[:, 0] + true_x = x.astype(jnp.float32) - assert jnp.allclose(pred_rot, true_rot) - assert jnp.allclose(pred_x, true_x) + assert jnp.allclose(pred_rot, true_rot, atol=1e-6, rtol=1e-6) + assert jnp.allclose(pred_x, true_x, atol=1e-6, rtol=1e-6) if __name__ == "__main__": - test_step() \ No newline at end of file + test_step() diff --git a/tests/test_initial_input_equinox.py b/tests/test_initial_input_equinox.py index 410025b..62ca64f 100644 --- a/tests/test_initial_input_equinox.py +++ b/tests/test_initial_input_equinox.py @@ -1,9 +1,10 @@ """Test all models on a simple 'remember the first input in the sequence' task""" -import pytest + import equinox as eqx import jax import jax.numpy as jnp import optax +import pytest from memax.equinox.train_utils import get_residual_memory_models @@ -29,9 +30,9 @@ def get_desired_accuracies(): "LinearRNN": 0.99, "PSpherical": 0.99, "GRU": 0.99, - "IndRNN": 0.55, - "Elman": 0.55, - "ElmanReLU": 0.55, + "IndRNN": 0.99, + "Elman": 0.60, + "ElmanReLU": 0.60, "Spherical": 0.99, "NMax": 0.99, "MGU": 0.99, @@ -44,11 +45,18 @@ def get_desired_accuracies(): def ce_loss(y_hat, y): return -jnp.mean(jnp.sum(y * jax.nn.log_softmax(y_hat, axis=-1), axis=-1)) -@pytest.mark.parametrize("model_name, model", get_residual_memory_models( - 4, 8, 4 - 1, key=jax.random.key(0), - ).items()) + +@pytest.mark.parametrize( + "model_name, model", + get_residual_memory_models( + 3, + 16, + 3 - 1, + key=jax.random.key(0), + ).items(), +) def test_initial_input( - model_name, model, epochs=2000, num_seqs=5, seq_len=20, input_dims=4 + model_name, model, epochs=400, num_seqs=5, seq_len=20, input_dims=3 ): timesteps = num_seqs * seq_len seq_idx = jnp.array([seq_len * i for i in range(num_seqs)]) @@ -111,7 +119,7 @@ def rerror(model, key): _, r_metrics = rerror(model, key) assert ( - r_metrics['accuracy']>= get_desired_accuracies()[model_name] + r_metrics["accuracy"] >= get_desired_accuracies()[model_name] ), f"Failed {model_name} (recurrent mode), expected {get_desired_accuracies()[model_name]}, got {r_metrics['accuracy']}" diff --git a/tests/test_initial_input_linen.py b/tests/test_initial_input_linen.py index 7e00390..13927f8 100644 --- a/tests/test_initial_input_linen.py +++ b/tests/test_initial_input_linen.py @@ -1,9 +1,11 @@ """Test all models on a simple 'remember the first input in the sequence' task""" -import pytest + +from functools import partial + import jax import jax.numpy as jnp import optax -from functools import partial +import pytest from memax.linen.train_utils import get_residual_memory_models @@ -16,14 +18,20 @@ def get_desired_accuracies(): "GRU": 0.999, } + def ce_loss(y_hat, y): return -jnp.mean(jnp.sum(y * jax.nn.log_softmax(y_hat, axis=-1), axis=-1)) -@pytest.mark.parametrize("model_name, model", get_residual_memory_models( - 8, 4 - 1, - ).items()) + +@pytest.mark.parametrize( + "model_name, model", + get_residual_memory_models( + 16, + 3 - 1, + ).items(), +) def test_initial_input( - model_name, model, epochs=4000, num_seqs=5, seq_len=20, input_dims=4 + model_name, model, epochs=400, num_seqs=5, seq_len=20, input_dims=3 ): timesteps = num_seqs * seq_len seq_idx = jnp.array([seq_len * i for i in range(num_seqs)]) @@ -34,7 +42,9 @@ def test_initial_input( key = jax.random.PRNGKey(0) dummy_x = jax.random.randint(key, (timesteps,), 0, input_dims - 1) dummy_x = jax.nn.one_hot(dummy_x, input_dims - 1) - dummy_x = jnp.concatenate([dummy_x, start.astype(jnp.float32).reshape(-1, 1)], axis=-1) + dummy_x = jnp.concatenate( + [dummy_x, start.astype(jnp.float32).reshape(-1, 1)], axis=-1 + ) dummy_h = model.zero_carry() dummy_starts = jnp.zeros(dummy_x.shape[0], dtype=bool) params = model.init(key, dummy_h, (dummy_x, dummy_starts)) @@ -43,7 +53,7 @@ def test_initial_input( state = opt.init(params) def error(params, key): - h = init_carry_fn(params) + h = init_carry_fn(params) x = jax.random.randint(key, (timesteps,), 0, input_dims - 1) x = jax.nn.one_hot(x, input_dims - 1) x = jnp.concatenate([x, start.astype(jnp.float32).reshape(-1, 1)], axis=-1) @@ -75,7 +85,7 @@ def error(params, key): # Verify recurrent mode works well too def rerror(params, key): - h = init_carry_fn(params) + h = init_carry_fn(params) x = jax.random.randint(key, (timesteps,), 0, input_dims - 1) x = jax.nn.one_hot(x, input_dims - 1) x = jnp.concatenate([x, start.astype(jnp.float32).reshape(-1, 1)], axis=-1) @@ -94,7 +104,7 @@ def rerror(params, key): _, r_metrics = rerror(params, key) assert ( - r_metrics['accuracy']>= get_desired_accuracies()[model_name] + r_metrics["accuracy"] >= get_desired_accuracies()[model_name] ), f"Failed {model_name} (recurrent mode), expected {get_desired_accuracies()[model_name]}, got {r_metrics['accuracy']}"