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

[simon-bachhuber]

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.

# 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) ?