How to optimize two set of parameters using one optimizer?

[GarfieldGa]

I have an encoder-decoder structure with two params. Currently, I tried to do something like below

grad_fn = jax.jit(jax.value_and_grad(self.calculate_loss, argnums=[1,2]), static_argnames=['self'])
loss, grads = grad_fn([self.encoder_params,self.decoder_params], x, y)
updates, self.opt_state = self.optimizer.update(grads, self.opt_state, self.decoder_params)
updated_params = optax.apply_updates([self.encoder_params,self.decoder_params], updates)

But this doesn't work and fails at the optax.apply_updates phase.

I wonder what a solution to this could be. The encoder I use is a simple CNN inherited from nn.Module, while the decoder is a customized class that is not a subclass of nn.module and decoder_params is just a predefined jnp.ndarray.

[jakevdp]

Can you put together a minimal reproducible example and add it to your question? It's hard to know what might be happening without seeing the code and its output. Thanks!