checkpoint

tests/test_checkpoint.py

@@ -0,0 +1,28 @@
+import runpy
+from pathlib import Path
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+
+SCRIPT_PATH = Path(__file__).parents[1] / "scripts" / "train.py"
+
+
+def test_initialize_network_restores_serialized_weights(tmp_path):
+    script = runpy.run_path(str(SCRIPT_PATH), run_name="train_script")
+    initialize_network = script["initialize_network"]
+
+    original = initialize_network(jrandom.PRNGKey(0))
+    checkpoint_path = tmp_path / "model.eqx"
+    eqx.tree_serialise_leaves(checkpoint_path, original)
+
+    restored = initialize_network(jrandom.PRNGKey(1), str(checkpoint_path))
+    original_leaves = jax.tree.leaves(eqx.filter(original, eqx.is_array))
+    restored_leaves = jax.tree.leaves(eqx.filter(restored, eqx.is_array))
+
+    assert len(original_leaves) == len(restored_leaves)
+    assert all(
+        jnp.array_equal(original_leaf, restored_leaf)
+        for original_leaf, restored_leaf in zip(original_leaves, restored_leaves, strict=True)
+    )