[state] Fixed a latent bug in `_call_primitive_discharge_rule`
The rule was passing a `WrappedFun` as an argument, which caused
```
jax._src.source_info_util.JaxStackTraceBeforeTransformation: TypeError: Error interpreting argument to named_computation as a JAX value. The problematic value is of type <class 'jax._src.linear_util.WrappedFun'> and was passed to named_computation at position 0.
```
PiperOrigin-RevId: 915954931
