JuicyTalk
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}")