Detailed Implementation on Equinox with Native JAX Modules, Filtered Transformations, Solid Layers, and End-to-End Training Flow

BATCH  = 128
EPOCHS = 30
steps_per_epoch = len(X_train) // BATCH
train_losses, val_losses = [], []


t0 = time.time()
for epoch in range(EPOCHS):
   key, sk = jax.random.split(key)
   perm = jax.random.permutation(sk, len(X_train))
   X_s, Y_s = X_train[perm], Y_train[perm]


   epoch_loss = 0.0
   for step in range(steps_per_epoch):
       xb = X_s[step*BATCH:(step+1)*BATCH]
       yb = Y_s[step*BATCH:(step+1)*BATCH]
       model, opt_state, loss = train_step(model, opt_state, xb, yb)
       epoch_loss += loss.item()


   val_loss = evaluate(model, X_val, Y_val).item()
   train_losses.append(epoch_loss / steps_per_epoch)
   val_losses.append(val_loss)


   if (epoch + 1) % 5 == 0:
       print(f"Epoch {epoch+1:3d}/{EPOCHS}  "
             f"train_loss={train_losses[-1]:.5f}  "
             f"val_loss={val_losses[-1]:.5f}")


print(f"nTotal training time: {time.time()-t0:.1f}s")


print("n" + "="*60)
print("SECTION 7: Save & load model weights")
print("="*60)


eqx.tree_serialise_leaves("model_weights.eqx", model)


key, mk2 = jax.random.split(key)
model_skeleton = ResNetMLP(1, 64, 1, n_blocks=4, key=mk2)
model_loaded   = eqx.tree_deserialise_leaves("model_weights.eqx", model_skeleton)


diff = jnp.max(jnp.abs(
   jax.tree_util.tree_leaves(eqx.filter(model, eqx.is_array))[0]
 - jax.tree_util.tree_leaves(eqx.filter(model_loaded, eqx.is_array))[0]
))
print(f"Max weight difference after reload: {diff:.2e}  (should be 0.0)")


fig, axes = plt.subplots(1, 2, figsize=(12, 4))


axes[0].plot(train_losses, label="Train MSE", color="#4C72B0")
axes[0].plot(val_losses,   label="Val MSE",   color="#DD8452", linestyle="--")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("MSE")
axes[0].set_title("Training curves")
axes[0].legend()
axes[0].grid(True, alpha=0.3)


x_plot  = jnp.linspace(-1, 1, 300).reshape(-1, 1)
y_true  = jnp.sin(2 * jnp.pi * x_plot)
y_pred  = jax.vmap(model)(x_plot)


axes[1].scatter(X_val[:100], Y_val[:100], s=10, alpha=0.4, color="gray", label="Data")
axes[1].plot(x_plot, y_true, color="#4C72B0",  linewidth=2, label="True f(x)")
axes[1].plot(x_plot, y_pred, color="#DD8452", linewidth=2, linestyle="--", label="Predicted")
axes[1].set_xlabel("x")
axes[1].set_ylabel("y")
axes[1].set_title("Sine regression fit")
axes[1].legend()
axes[1].grid(True, alpha=0.3)


plt.tight_layout()
plt.savefig("equinox_tutorial.png", dpi=150)
plt.show()
print("nDone! Plot saved to equinox_tutorial.png")


print("n" + "="*60)
print("BONUS: eqx.filter_jit + shape inference debug tip")
print("="*60)


jaxpr = jax.make_jaxpr(jax.vmap(model))(x_plot)
n_eqns = len(jaxpr.jaxpr.eqns)
print(f"Compiled ResNetMLP jaxpr has {n_eqns} equations (ops) for batch input {x_plot.shape}")
BATCH  = 128
EPOCHS = 30
steps_per_epoch = len(X_train) // BATCH
train_losses, val_losses = [], []


t0 = time.time()
for epoch in range(EPOCHS):
   key, sk = jax.random.split(key)
   perm = jax.random.permutation(sk, len(X_train))
   X_s, Y_s = X_train[perm], Y_train[perm]


   epoch_loss = 0.0
   for step in range(steps_per_epoch):
       xb = X_s[step*BATCH:(step+1)*BATCH]
       yb = Y_s[step*BATCH:(step+1)*BATCH]
       model, opt_state, loss = train_step(model, opt_state, xb, yb)
       epoch_loss += loss.item()


   val_loss = evaluate(model, X_val, Y_val).item()
   train_losses.append(epoch_loss / steps_per_epoch)
   val_losses.append(val_loss)


   if (epoch + 1) % 5 == 0:
       print(f"Epoch {epoch+1:3d}/{EPOCHS}  "
             f"train_loss={train_losses[-1]:.5f}  "
             f"val_loss={val_losses[-1]:.5f}")


print(f"nTotal training time: {time.time()-t0:.1f}s")


print("n" + "="*60)
print("SECTION 7: Save & load model weights")
print("="*60)


eqx.tree_serialise_leaves("model_weights.eqx", model)


key, mk2 = jax.random.split(key)
model_skeleton = ResNetMLP(1, 64, 1, n_blocks=4, key=mk2)
model_loaded   = eqx.tree_deserialise_leaves("model_weights.eqx", model_skeleton)


diff = jnp.max(jnp.abs(
   jax.tree_util.tree_leaves(eqx.filter(model, eqx.is_array))[0]
 - jax.tree_util.tree_leaves(eqx.filter(model_loaded, eqx.is_array))[0]
))
print(f"Max weight difference after reload: {diff:.2e}  (should be 0.0)")


fig, axes = plt.subplots(1, 2, figsize=(12, 4))


axes[0].plot(train_losses, label="Train MSE", color="#4C72B0")
axes[0].plot(val_losses,   label="Val MSE",   color="#DD8452", linestyle="--")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("MSE")
axes[0].set_title("Training curves")
axes[0].legend()
axes[0].grid(True, alpha=0.3)


x_plot  = jnp.linspace(-1, 1, 300).reshape(-1, 1)
y_true  = jnp.sin(2 * jnp.pi * x_plot)
y_pred  = jax.vmap(model)(x_plot)


axes[1].scatter(X_val[:100], Y_val[:100], s=10, alpha=0.4, color="gray", label="Data")
axes[1].plot(x_plot, y_true, color="#4C72B0",  linewidth=2, label="True f(x)")
axes[1].plot(x_plot, y_pred, color="#DD8452", linewidth=2, linestyle="--", label="Predicted")
axes[1].set_xlabel("x")
axes[1].set_ylabel("y")
axes[1].set_title("Sine regression fit")
axes[1].legend()
axes[1].grid(True, alpha=0.3)


plt.tight_layout()
plt.savefig("equinox_tutorial.png", dpi=150)
plt.show()
print("nDone! Plot saved to equinox_tutorial.png")


print("n" + "="*60)
print("BONUS: eqx.filter_jit + shape inference debug tip")
print("="*60)


jaxpr = jax.make_jaxpr(jax.vmap(model))(x_plot)
n_eqns = len(jaxpr.jaxpr.eqns)
print(f"Compiled ResNetMLP jaxpr has {n_eqns} equations (ops) for batch input {x_plot.shape}")