clipping model parameters in flax.training.train_state.TrainState

[jecampagne]

Hi,
I'm quite new in flax and I'm using linen nn.Module to setup a simple model. Then I'm using flax.training.train_state.TrainState default to create a state as followed

@partial(jax.pmap, static_broadcasted_argnums=(1, 2))
def create_state(rng, model_cls, input_shape): 
  r"""Create the training state given a model class. """ 
  model = model_cls()

  tx = optax.rmsprop(5e-5)
  variables = model.init(rng, jnp.ones(input_shape))

  state = TrainState.create(apply_fn=model.apply, tx=tx, 
      params=variables['params'])
  
  return state

This state once created is used in a training loop as

  new_state = model_state.apply_gradients(grads=grads)

Now, I wander how I can implement a variables['params'] clipping in this context.
Thanks for your help.