jax → ONNX export via jax2onnx: jaxtrain now produces HSSM-consumable artifacts - #109
Open
AlexanderFengler wants to merge 1 commit into
Open
jax → ONNX export via jax2onnx: jaxtrain now produces HSSM-consumable artifacts#109AlexanderFengler wants to merge 1 commit into
AlexanderFengler wants to merge 1 commit into
Conversation
…ifacts jaxtrain was a dead end for the ecosystem: it saved flax parameter bytes that nothing downstream could read — transform-onnx only handles torch state dicts and HSSM consumes ONNX exclusively. This closes the single missing edge. - New lanfactory/onnx/jax_export.py: export_forward_to_onnx (shared core) + transform_jax_to_onnx (file-based, for retroactive conversion of existing .jax artifacts), exposed as the `transform-jax-onnx` CLI - The jax trainer exports ONNX alongside its other artifacts (mirroring the torch trainer's _save_onnx) and logs it to MLflow; `--no-export-onnx` on jaxtrain opts out - Dependency: jax2onnx==0.15.* (exact-pinned 0.x line), opset 17 to match the sbi/bayesflow exporters Two contract decisions, both verified empirically: - The graph is traced with a concrete (1, input_dim) dummy — every dim static, zero dynamic axes — matching the torch MLP exporter and the production networks on franklab/HSSM, NOT rank-1: jax2onnx lowers nn.Dense to Gemm, whose ONNX spec requires rank-2 inputs (a rank-1 trace is rejected by onnxruntime). Verified end-to-end that HSSM's make_jax_func loads the export and rank-1-per-trial + jax.vmap consumption matches exactly. - The exported graph is the EVAL head: identity for logprob (LAN), logsigmoid for logits (CPN/OPN) — matching every torch export path (_save_onnx calls .eval(); transform-onnx exports under torch's EVAL default) and HSSM's element-wise log-likelihood consumption. Exporting the raw training head would silently corrupt every CPN/OPN logp by +log(1+exp(-logit)); a test asserts the logits export equals the logsigmoid head and is NOT the raw head. Tests: contract (all dims concrete, (1,D) shape, op allowlist calibrated against jax2onnx 0.15's actual lowering — relu emits Max, the logsigmoid head emits Neg/Exp/Add/Log), parity (ORT vs eval-head jax forward over 1000 draws, atol 1e-4), file-based vs in-trainer export equivalence, logits-head correctness, and the opt-out flag. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
PR-0.3, stacked on #108 (review only the top commit; retarget to
mainafter #108 merges). Closes the gap that madejaxtraina dead end for HSSM: the jax trainer saved flax parameter bytes that nothing downstream could convert —transform-onnxonly reads torch state dicts, and HSSM consumes ONNX exclusively.lanfactory/onnx/jax_export.py(~150 lines incl. docs): wrapsjax2onnx(pinned==0.15.*) over the existingJaxMLPFactory.make_forward_partialrestore path — minimal-carry, no bespoke tracertransform-jax-onnxCLI (file-based:network_config.pickle+*_train_state.jax→.onnx)jaxtraingains--export-onnx/--no-export-onnx(default on): the trainer exports ONNX right after training and logs it to MLflow(1, input_dim)dummy, no dynamic axes.(1, D)rather than rank-1 because jax2onnx lowersnn.DensetoGemm, whose ONNX spec requires rank-2 (the sbi/bayesflow exporters trace rank-1 because torch lowers Linear to rank-agnostic MatMul+Add — different tracer, different constraint). Verified end-to-end against HSSM'smake_jax_func+jax.vmapconsumption with exact paritylogprob(LAN) this equals the raw head; forlogits(CPN/OPN) it applies logsigmoid, matching every torch export path. Exporting the raw training head would silently corrupt every downstream logp by+log(1+exp(-logit))— this was caught in review and is now explicitly regression-testedtransform-onnxuntouched — they remain the verified fallback pathTests (the parity test is the load-bearing one)
(1, D), op profile ⊆ {Gemm + elementwise} (guards against jax2onnx lowering changes under the 0.x pin)atol=1e-4; file-based transform ≡ in-trainer export-log1p(exp(-raw))and provably ≠ the raw head--no-export-onnxproduces no.onnxCommands run
🤖 Generated with Claude Code