Correctness checking vs existing PyTorch implementation?

[mdda]

Seems to me that it would make sense to have :

• some kind of fuzzer, where inputs of the right shapes are created randomly, and applied to both implementations; and/or
• an automatic utility that can put function-wise or line-wise logging on existing implementations, so that internal states could be measured for a given layer in-situ, when 'real inputs' are run through the full model.

Extension : Armed with these verification tools, we could get an LLM (with verification in the loop) to do the translation for us (given a small library of working PyTorch/JAX translation pairs, which would easily fit in Gemini context).

[monatis]

Wow this would be really awesome --agentic AI engineer :-)

Actually, transformers models already return intermediate states if requested, and on the Jax/Flax side, we can make use of Module.sow() functionality to get intermediates. We may also need to put some Flax docs in the Gemini context because flax.nnx is a newer API and LLMs are still only aware of the flax.linen API.

Sounds quite reasonable to me.

[monatis]

I played with this idea a little bit but not only Gemini but also other LLMs seem to be too bad in recent Flax versions, especially NNX. Maybe it might be a prototype if someone wants to put some effort to it, but I myself won't continue with this idea.