Updated demo to work with newer version of JAX

Since new versions of JAX don't support the abstraction using `__jax_array__` anymore, the old method of passing `keras.Variables` to `jax.jit` compiled functions doesn't work.

This change fixes that by manually extracting the underlying jax arrays

Author: vinayak19th

guides/writing_a_custom_training_loop_in_jax.py

@@ -277,9 +277,9 @@ def train_step(state, data):
 # Build optimizer variables.
 optimizer.build(model.trainable_variables)
 
-trainable_variables = model.trainable_variables
-non_trainable_variables = model.non_trainable_variables
-optimizer_variables = optimizer.variables
+trainable_variables = [v.value for v in model.trainable_variables]
+non_trainable_variables = [v.value for v in model.non_trainable_variables]
+optimizer_variables = [v.value for v in optimizer.variables]
 state = trainable_variables, non_trainable_variables, optimizer_variables
 
 # Training loop