Question about functional training loop documentation

[vladfi1]

In https://flax.readthedocs.io/en/stable/guides/performance.html#functional-training-loop I see that the graphdef itself is passed through the jitted jax_train_step. This seems odd to me as the graphdef presumably doesn't have any jax.Arrays. Would it make more sense to simply access the outer graphdef instead?