jax

fa827a59 - [jax2tf] Added the last comments from the jax2tf doc inside the (#4249)

[jax2tf] Added the last comments from the jax2tf doc inside the (#4249) correctness_stats code. In principle, all the relevant documentation that was in the doc has been moved to the new documentation & comments of categorize.