How to pass a Flax dataclass to a JITed function where some fields are static and some are dynamic?

[marcvanzee]

Original question by @malmaud.

[marcvanzee]

This works:

class Params(flax.struct.PyTreeNode):
  a: bool = flax.struct.field(pytree_node=False)  # static
  b: int  # dynamic

Note that you can also use @flax.struct.dataclass to decorate Params rather than subclassing flax.struct.PyTreeNode, but the later currently gives better PyType support. PyType sometimes complains about either not understanding the constructor or the replace method when using the annotation.