optimizer = optax.adam(learning_rate=0.001)