GenCast training

[dkokron]

I am playing with the gencast_mini_demo.ipynb using the 'GenCast 1p0deg Mini <2019.npz' model along with the source-era5_date-2019-03-29_res-1.0_levels-13_steps-01.nc dataset. I added the following commands after the "Gradient computation" cell to see if the training steps changed anything in the params structure. Nothing changed. So the question becomes, what is the training doing? Which fields should I be comparing?

print(np.array_equal(ckpt.params["fourier_features_mlp/~/mlp/~/linear_0"]["b"], params["fourier_features_mlp/~/mlp/~/linear_0"]["b"]))
print(np.array_equal(ckpt.params["fourier_features_mlp/~/mlp/~/linear_0"]["w"], params["fourier_features_mlp/~/mlp/~/linear_0"]["w"]))
print(np.array_equal(ckpt.params["fourier_features_mlp/~/mlp/~/linear_1"]["b"], params["fourier_features_mlp/~/mlp/~/linear_1"]["b"]))
print(np.array_equal(ckpt.params["fourier_features_mlp/~/mlp/~/linear_1"]["w"], params["fourier_features_mlp/~/mlp/~/linear_1"]["w"])) 

True
True
True
True

[alvarosg]

The grads_fn_jitted only produces the gradients of the params. In order to actually perform an update of the params, then you need to use an optimizer (e.g. Adam from the optax library) to used those gradients to compute an update to the params.