Start work on proper experiments

README.md

@@ -45,15 +45,15 @@ We provide [datasets](memax/datasets) to test our recurrent models.
 >
 > **Sequence Lengths:** `[784]`
 
-### MNIST Math [[HuggingFace]](https://huggingface.co/datasets?sort=trending&search=bolt-lab%2Fmnist-math) [[Code]](memax/datasets/sequential_mnist.py)
+### MNIST Math [[HuggingFace]](https://huggingface.co/datasets/bolt-lab/mnist-math-100) [[Code]](memax/datasets/mnist_math.py)
 > The recurrent model receives a sequence of MNIST images and operators, pixel by pixel, and must predict the percentile of the operators applied to the MNIST image classes.
 >
-> **Sequence Lengths:** `[784 * 5, 784 * 100, 784 * 1_000, 784 * 10_000, 784 * 1_000_000]`
+> **Sequence Lengths:** `[784 * 100, …]` on Hub as `bolt-lab/mnist-math-{length}` (default loader: `100`; `100_000` is under `smorad/mnist-math-100000`).
 
-### Continuous Localization [[HuggingFace]](https://huggingface.co/datasets?sort=trending&search=bolt-lab%2Fcontinuous-localization) [[Code]](memax/datasets/sequential_mnist.py)
+### Continuous Localization [[HuggingFace]](https://huggingface.co/datasets/bolt-lab/continuous-localization-20) [[Code]](memax/datasets/continuous_localization.py)
 > The recurrent model receives a sequence of translation and rotation vectors **in the local coordinate frame**, and must predict the corresponding position and orientation **in the global coordinate frame**.
 >
-> **Sequence Lengths:** `[20, 100, 1_000]`
+> **Sequence Lengths:** `[20, 100]` on Hub as `bolt-lab/continuous-localization-{length}`.
 
 # Getting Started
 Install `memax` using pip for your specific framework:
@@ -105,10 +105,23 @@ hs, ys = filter_jit(filter_vmap(model))(hs_0, (xs, starts))
 ```
 
 ## Running Baselines
-You can compare various recurrent models on our datasets with a single command
+You can compare various recurrent models on our datasets with a single command:
 ```bash
-python run_equinox_experiments.py # equinox framework
-python run_linen_experiments.py # flax linen framework
+python run_equinox_experiments.py --dataset-name sequential_mnist
+python run_linen_experiments.py --dataset-name sequential_mnist
+```
+
+Dataset names: see `memax.experiments.datasets.DATASET_NAMES` (e.g. `sequential_mnist`, `mnist_math`, `mnist-math-100`, `continuous-localization-20`). Legacy alias: `sequential_rotation`.
+
+For a quick CPU smoke run (tiny model, few optimizer steps):
+```bash
+python run_equinox_experiments.py --dataset-name sequential_mnist --smoke --models GRU
+python run_linen_experiments.py --dataset-name sequential_mnist --smoke --models GRU
+```
+
+CI runs synthetic smoke tests only (`pytest`; no Hugging Face downloads). To exercise real datasets locally:
+```bash
+pytest -m integration tests/test_experiment_smoke.py
 ```
 
 

memax/datasets/continuous_localization.py

@@ -1,37 +1,46 @@
 """This module generates the continuous localization dataset and uploads it to Hugging Face.
 The dataset consists of sequences of 3D rotations and translations, along with the absolute positions and
-orientations. Each sequence is generated by applying a series of random small rotations and translations to an initial position and orientation. The goal is to predict the absolute position and orientation from the sequence of relative movements."""
+orientations. Each sequence is generated by applying a series of random small rotations and translations to an initial position and orientation. The goal is to predict the absolute position and orientation from the sequence of relative movements.
+"""
 
-import jax.numpy as jnp
 import jax
+import jax.numpy as jnp
+from datasets import Array2D, Dataset, Features, load_dataset
 from jax.scipy.spatial.transform import Rotation
-from datasets import Dataset, Features, Array2D, load_dataset
 
-SEQ_LENS = [100, 1_000, 10_000, 100_000, 1_000_000]
+from memax.datasets.hub import (
+    CONTINUOUS_LOCALIZATION_ORG,
+    CONTINUOUS_LOCALIZATION_SEQ_LENS,
+    continuous_localization_hub_id,
+)
+
+# Lengths used when generating/uploading new Hub revisions.
+UPLOAD_SEQ_LENS = [100, 1_000, 10_000, 100_000, 1_000_000]
+
 
 def step(carry, inputs):
-    (x, rot) = carry
-    (dx, drot) = inputs
+    x, rot = carry
+    dx, drot = inputs
     transform_x = jax.vmap(lambda rot_t, x_t, dx: rot_t.apply(dx) + x_t)
     transform_r = jax.vmap(lambda rot_t, drot: rot_t * drot)
     x = transform_x(rot, x, dx)
     rot = transform_r(rot, drot)
     return ((x, rot), (x, rot))
 
 
-def generate_dataset(
-        key,
-        batch_size,
-        num_steps
-):
+def generate_dataset(key, batch_size, num_steps):
     spatial_dim = 3
     keys = jax.random.split(key, 4)
     x = jnp.zeros((batch_size, spatial_dim))
     rot = jax.vmap(Rotation.identity, axis_size=batch_size)()
     dx_mag = jax.random.exponential(keys[0], (num_steps, batch_size, 1)) * 2
-    dx_dir = jax.random.uniform(keys[1], (num_steps, batch_size, 3), minval=-1., maxval=1.)
+    dx_dir = jax.random.uniform(
+        keys[1], (num_steps, batch_size, 3), minval=-1.0, maxval=1.0
+    )
     dx = dx_dir / jnp.linalg.norm(dx_dir, axis=-1, keepdims=True) * dx_mag
-    drot_dir = jax.random.uniform(keys[2], (num_steps, batch_size, spatial_dim), minval=-1, maxval=1)
+    drot_dir = jax.random.uniform(
+        keys[2], (num_steps, batch_size, spatial_dim), minval=-1, maxval=1
+    )
     drot_mag = jax.random.uniform(keys[3], (num_steps, batch_size, 1), maxval=jnp.pi)
     drot_vec = drot_dir / jnp.linalg.norm(drot_dir, axis=-1, keepdims=True) * drot_mag
     drot = Rotation.from_rotvec(drot_vec)
@@ -45,93 +54,97 @@ def generate_dataset(
     outputs = jnp.concatenate([rot_abs_vec, x_abs], axis=-1)
 
     return {
-        "inputs": jnp.permute_dims(inputs, (1,0,2)),
-        "outputs": jnp.permute_dims(outputs, (1,0,2)),
-        "delta_rotation": jnp.permute_dims(drot_vec, (1,0,2)),
-        "delta_position": jnp.permute_dims(dx, (1,0,2)),
-        "absolute_rotation": jnp.permute_dims(rot_abs_vec, (1,0,2)),
-        "absolute_position": jnp.permute_dims(x_abs, (1,0,2)),
+        "inputs": jnp.permute_dims(inputs, (1, 0, 2)),
+        "outputs": jnp.permute_dims(outputs, (1, 0, 2)),
+        "delta_rotation": jnp.permute_dims(drot_vec, (1, 0, 2)),
+        "delta_position": jnp.permute_dims(dx, (1, 0, 2)),
+        "absolute_rotation": jnp.permute_dims(rot_abs_vec, (1, 0, 2)),
+        "absolute_position": jnp.permute_dims(x_abs, (1, 0, 2)),
     }
 
 
 def upload_hf_datasets(
-    batch_size_train = 1e6,
-    batch_size_test = 2e5,
-    sequence_length = 20,
+    batch_size_train=1e6,
+    batch_size_test=2e5,
+    sequence_length=20,
 ):
-    name = f"bolt-lab/continuous-localization-{sequence_length}",
+    name = f"{CONTINUOUS_LOCALIZATION_ORG}/continuous-localization-{sequence_length}"
     batch_size_train = int(batch_size_train)
     batch_size_test = int(batch_size_test)
-    FEATURES = Features({
-        "inputs": Array2D(dtype='float32', shape=(sequence_length, 6)),
-        "outputs": Array2D(dtype='float32', shape=(sequence_length, 6)),
-        "delta_rotation": Array2D(dtype='float32', shape=(sequence_length, 3)),
-        "delta_position": Array2D(dtype='float32', shape=(sequence_length, 3)),
-        "absolute_rotation": Array2D(dtype='float32', shape=(sequence_length, 3)),
-        "absolute_position": Array2D(dtype='float32', shape=(sequence_length, 3)),
-    })
+    FEATURES = Features(
+        {
+            "inputs": Array2D(dtype="float32", shape=(sequence_length, 6)),
+            "outputs": Array2D(dtype="float32", shape=(sequence_length, 6)),
+            "delta_rotation": Array2D(dtype="float32", shape=(sequence_length, 3)),
+            "delta_position": Array2D(dtype="float32", shape=(sequence_length, 3)),
+            "absolute_rotation": Array2D(dtype="float32", shape=(sequence_length, 3)),
+            "absolute_position": Array2D(dtype="float32", shape=(sequence_length, 3)),
+        }
+    )
     key = jax.random.key(0)
     keys = jax.random.split(key, 2)
-    train_dict = generate_dataset(keys[0], batch_size=batch_size_train, num_steps=sequence_length)
-    test_dict = generate_dataset(keys[1], batch_size=batch_size_test, num_steps=sequence_length)
-    train_dataset = Dataset.from_dict(train_dict).cast(features=FEATURES, batch_size=batch_size_train)
-    test_dataset = Dataset.from_dict(test_dict).cast(features=FEATURES, batch_size=batch_size_test)
+    train_dict = generate_dataset(
+        keys[0], batch_size=batch_size_train, num_steps=sequence_length
+    )
+    test_dict = generate_dataset(
+        keys[1], batch_size=batch_size_test, num_steps=sequence_length
+    )
+    train_dataset = Dataset.from_dict(train_dict).cast(
+        features=FEATURES, batch_size=batch_size_train
+    )
+    test_dataset = Dataset.from_dict(test_dict).cast(
+        features=FEATURES, batch_size=batch_size_test
+    )
     train_dataset.push_to_hub(name, split="train")
     test_dataset.push_to_hub(name, split="test")
 
+
 def get_rot_dataset(sequence_length=20):
-    seq_lens = [20]
-    assert sequence_length in seq_lens, f"Invalid sequenec length, must be one of {seq_lens}"
-    num_labels = 10
-    dataset = load_dataset(f"bolt-lab/continuous-localization-{sequence_length}")
+    assert (
+        sequence_length in CONTINUOUS_LOCALIZATION_SEQ_LENS
+    ), f"Invalid sequence length, must be one of {CONTINUOUS_LOCALIZATION_SEQ_LENS}"
+    dataset = load_dataset(continuous_localization_hub_id(sequence_length))
 
     x = jnp.array(dataset["train"]["delta_rotation"])
-    y = jnp.array(dataset["train"]["absolute_rotation"][-1])
-    #y = jnp.array(dataset["train"]["absolute_rotation_percentile"])
-    # TODO: One-hot based on abs of rotation angle?
-    #y = jax.nn.one_hot(y, num_labels)
+    y = jnp.array(dataset["train"]["absolute_rotation"])[:, -1, :]
 
     test_x = jnp.array(dataset["test"]["delta_rotation"])
-    #test_y = jnp.array(dataset["test"]["absolute_rotation_percentile"])
-    test_y = jnp.array(dataset["test"]["absolute_rotation"][-1])
-    #test_y = jax.nn.one_hot(test_y, num_labels)
+    test_y = jnp.array(dataset["test"]["absolute_rotation"])[:, -1, :]
 
     return {
         "x_train": x,
         "y_train": y,
         "x_test": test_x,
         "y_test": test_y,
         "size": x.shape[0],
-        "num_labels": 10,
+        "num_labels": y.shape[-1],
     }
 
+
 def get_trans_dataset(sequence_length=20):
-    seq_lens = [20, 1024]
-    assert sequence_length in seq_lens, f"Invalid sequenec length, must be one of {seq_lens}"
-    dataset = load_dataset(f"bolt-lab/continuous-localization-{sequence_length}")
-
-    x = jnp.concatenate([
-        jnp.array(dataset["train"]["delta_rotation"]),
-        jnp.array(dataset["train"]["delta_translation"])
-    ], axis=-1)
-    y = jnp.array(dataset["train"]["absolute_translation_percentile"])
-
-    test_x = jnp.concatenate([
-        jnp.array(dataset["test"]["delta_rotation"]),
-        jnp.array(dataset["test"]["delta_translation"])
-    ], axis=-1)
-    test_y = jnp.array(dataset["test"]["absolute_translation_percentile"])
+    assert (
+        sequence_length in CONTINUOUS_LOCALIZATION_SEQ_LENS
+    ), f"Invalid sequence length, must be one of {CONTINUOUS_LOCALIZATION_SEQ_LENS}"
+    dataset = load_dataset(continuous_localization_hub_id(sequence_length))
+
+    x = jnp.array(dataset["train"]["inputs"])
+    y = jnp.array(dataset["train"]["outputs"])[:, -1, :]
+
+    test_x = jnp.array(dataset["test"]["inputs"])
+    test_y = jnp.array(dataset["test"]["outputs"])[:, -1, :]
 
     return {
         "x_train": x,
         "y_train": y,
         "x_test": test_x,
         "y_test": test_y,
         "size": x.shape[0],
-        "num_labels": 10,
+        "num_labels": y.shape[-1],
     }
 
 
 if __name__ == "__main__":
-    for seq_len in SEQ_LENS:
-        upload_hf_datasets(batch_size_train=60_000, batch_size_test=10_000, sequence_length=seq_len)
+    for seq_len in UPLOAD_SEQ_LENS:
+        upload_hf_datasets(
+            batch_size_train=60_000, batch_size_test=10_000, sequence_length=seq_len
+        )

memax/datasets/hub.py

@@ -0,0 +1,47 @@
+"""Hugging Face Hub dataset identifiers used by memax loaders.
+
+Canonical dataset pages:
+- https://huggingface.co/datasets/ylecun/mnist
+- https://huggingface.co/datasets/bolt-lab/mnist-math-{seq_len}
+- https://huggingface.co/datasets/bolt-lab/continuous-localization-{seq_len}
+"""
+
+# Raw MNIST images (README: ylecun/mnist). The legacy id "mnist" redirects but logs hub warnings.
+MNIST = "ylecun/mnist"
+
+# Published under bolt-lab; 100_000 lives only under smorad.
+MNIST_MATH_ORG = "bolt-lab"
+MNIST_MATH_ORG_ALT = "smorad"
+MNIST_MATH_SEQ_LENS = [100, 1_000, 10_000, 100_000, 1_000_000]
+
+CONTINUOUS_LOCALIZATION_ORG = "bolt-lab"
+# Datasets available on the Hub as of upload (longer seq_lens may be added separately).
+CONTINUOUS_LOCALIZATION_SEQ_LENS = [20, 100]
+
+
+def mnist_math_hub_id(seq_len: int) -> str:
+    """Return the Hub repo id for a MNIST Math sequence length."""
+    if seq_len not in MNIST_MATH_SEQ_LENS:
+        raise ValueError(
+            f"Invalid seq_len {seq_len}, must be one of {MNIST_MATH_SEQ_LENS}"
+        )
+    if seq_len == 100_000:
+        return f"{MNIST_MATH_ORG_ALT}/mnist-math-{seq_len}"
+    return f"{MNIST_MATH_ORG}/mnist-math-{seq_len}"
+
+
+def mnist_math_hub_url(seq_len: int) -> str:
+    return f"https://huggingface.co/datasets/{mnist_math_hub_id(seq_len)}"
+
+
+def continuous_localization_hub_id(sequence_length: int) -> str:
+    if sequence_length not in CONTINUOUS_LOCALIZATION_SEQ_LENS:
+        raise ValueError(
+            f"Invalid sequence_length {sequence_length}, "
+            f"must be one of {CONTINUOUS_LOCALIZATION_SEQ_LENS}"
+        )
+    return f"{CONTINUOUS_LOCALIZATION_ORG}/continuous-localization-{sequence_length}"
+
+
+def continuous_localization_hub_url(sequence_length: int) -> str:
+    return f"https://huggingface.co/datasets/{continuous_localization_hub_id(sequence_length)}"

memax/datasets/mnist_listops.py

@@ -4,9 +4,9 @@
 import jax.numpy as jnp
 from datasets import load_dataset  # huggingface datasets
 
+from memax.datasets.hub import MNIST
 from memax.train_utils import get_residual_memory_models
 
-
 NUM_EPOCHS = 100
 BATCH_SIZE = 32
 SEQ_LEN = 784
@@ -105,7 +105,7 @@ def normalize_and_flatten(x):
 
 
 def make_dataset(dataset_size=3, num_terms=5, key=jax.random.key(0), batch_size=32):
-    dataset = load_dataset("mnist")
+    dataset = load_dataset(MNIST)
     x, y = jnp.array(dataset["train"]["image"]), jnp.array(dataset["train"]["label"])
 
     keys = jax.random.split(key, 3)

memax/datasets/mnist_math.py

@@ -1,6 +1,6 @@
 """
 This file contains the code for creating the MNIST Math dataset,
-as well as the code that preprocesses the dataset before training. 
+as well as the code that preprocesses the dataset before training.
 The MNIST Math dataset consists of sequences of MNIST digits
 interspersed with plus and minus operators. The task is to compute
 the cumulative result of the arithmetic expression formed by the digits
@@ -16,8 +16,10 @@
 import tqdm
 from datasets import Dataset, Features, Image, Sequence, Value, load_dataset
 
+from memax.datasets.hub import MNIST, MNIST_MATH_SEQ_LENS, mnist_math_hub_id
+
 NUM_LABELS = 10
-SEQ_LENS = [100, 1_000, 10_000, 100_000, 1_000_000]
+SEQ_LENS = MNIST_MATH_SEQ_LENS
 
 FEATURES = Features(
     {
@@ -173,7 +175,7 @@ def generate_dataset(key, mnist, dataset_size=10, num_terms=5):
 
 
 def make_hf_datasets():
-    mnist = load_dataset("mnist")
+    mnist = load_dataset(MNIST)
     train_dset = generate_dataset(jax.random.key(0), mnist["train"], 60_000)
     test_dset = generate_dataset(jax.random.key(1), mnist["test"], 10_000)
     train_dset = Dataset.from_dict(train_dset).cast(FEATURES)
@@ -183,13 +185,14 @@ def make_hf_datasets():
 
 def upload_hf_datasets(length):
     train, test = make_hf_datasets()
-    train.push_to_hub(f"smorad/mnist-math-{length}", split="train")
-    test.push_to_hub(f"smorad/mnist-math-{length}", split="test")
+    repo_id = mnist_math_hub_id(length)
+    train.push_to_hub(repo_id, split="train")
+    test.push_to_hub(repo_id, split="test")
 
 
-def get_dataset(seq_len=5):
+def get_dataset(seq_len=100):
     assert seq_len in SEQ_LENS, f"Invalid length, must be in {SEQ_LENS}"
-    dataset = load_dataset("smorad/mnist-math").with_format("np")
+    dataset = load_dataset(mnist_math_hub_id(seq_len)).with_format("np")
     num_labels = 10
 
     x = dataset["train"]["image_flat"]

memax/datasets/sequential_mnist.py

@@ -7,6 +7,8 @@
 import jax.numpy as jnp
 from datasets import load_dataset
 
+from memax.datasets.hub import MNIST
+
 
 def normalize_and_flatten(x):
     # batch, time, feature
@@ -15,7 +17,7 @@ def normalize_and_flatten(x):
 
 
 def get_dataset():
-    dataset = load_dataset("mnist")
+    dataset = load_dataset(MNIST)
     num_labels = 10
 
     x = jnp.array(dataset["train"]["image"])
@@ -35,4 +37,4 @@ def get_dataset():
         "y_test": test_y,
         "num_labels": num_labels,
         "size": x.shape[0],
-    }
+    }