Specify Module's PyTree-Representation for jit/grad seperately. I.e. How to freeze state.variables

Disclaimer: I have not used `oryx` yet. Further, not an issue but rather just a question/discussion.

Suppose i want to define some recurrent network but its initial hidden state is *not* a parameter, i.e. it should be exposed to `jax.jit` but not to `jax.grad`. How can this be done?

E.g.
```python
# syntax might be slightly wrong, think of it as pseudo-code
def network_def(x):
s = state.variable(..., name="hidden-state")
p = state.variable(..., name="parameters")
s, y = f(s, p, x)
state.assign(s, name="hidden-state")
return y

network = state.init(network_def)(x)

@jax.jit # <- this should "see" hidden-state
@jax.grad # <- this should not "see" hidden-state
def loss_fn(network, x, y):
...
```

Is there an elegant way of doing that?
Thank you!

Also, are all jax-transformations supported? Readme mentions `jit`, `grad`, `vmap`. What about `pmap`,`scan` (and all the others) ?

Start by reading the README sections describing state.variable, state.assign, state.init, and the supported JAX transformations. Determine whether the issue calls for documenting existing behavior or designing separate PyTree handling for jit and grad, including the status of pmap and scan.

Written by the indexing model from the issue text.