28 lines
995 B
Python
28 lines
995 B
Python
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)
|
|
)
|