Repository navigation
fix(export): correct and complete the PyTorch and JAX operator mappings - #1415
Merged
Merged
Conversation
Wrong results: - JAX mapped `sympy.LessThan` (`<=`) to `jnp.less` and `GreaterThan` (`>=`) to `jnp.greater`, so `less_equal`/`greater_equal` were wrong at ties; the strict relations had no mapping at all. - JAX mapped `Max`/`Min` to the reductions `jnp.max`/`jnp.min`, so the second argument was read as an axis: `max(x0, 0.5)` raised a `TypeError`. Both backends also failed on `Max`/`Min`/`And`/`Or` with more than two arguments, which SymPy produces from nested calls such as `max(max(x0, x1), x2)`. These are now elementwise and folded pairwise. - PyTorch looked up built-in mappings before `extra_torch_mappings`, so users could not override them (the JAX export already let them). - PyTorch shared one trainable parameter between equal constants, e.g. `x0 * 2.5 + x1 * 2.5` had one parameter, versus two in the search and in the JAX export. Missing support: - `Piecewise`, which the SymPy forms of `greater`, `less`, `greater_equal`, `less_equal`, `cond`, `logical_or`, `logical_and`, `relu` and `clamp` use, had no PyTorch or JAX mapping, so exporting any of these operators raised a `KeyError`. Map it to nested `torch.where`/`jnp.where`, with NaN where no branch applies, as in SymPy. This finishes the work started in #433. - Add `erfc` for PyTorch and `gamma`/`loggamma` for JAX. Robustness: - `pytorch()` and `jax()` set `output_torch_format`/`output_jax_format` before refreshing. When one equation could not be exported, the flag stayed set, so every later `refresh()`, and therefore `sympy()` and `latex()`, raised the same error. Restore the flag when the export fails. - `torch_initialized`/`jax_initialized` were never set, so the torch module classes were redefined on every export. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Both exporters now translate Piecewise (used by greater, less, cond, relu, clamp, logical_and/or) into nested where calls, and fold Max/Min over any number of arguments. JAX maps Max/Min to the elementwise jnp.maximum/jnp.minimum and maps strict and non-strict relations correctly. User extra_torch_mappings now override built-in mappings, and both exporters cache their lazy imports. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
MilesCranmer
reviewed
Oct 8, 2026
Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
MilesCranmer
reviewed
Oct 8, 2026
MilesCranmer
approved these changes
Oct 8, 2026
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
3 participants
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.
Summary
Fixes PyTorch and JAX export, finishing the
Piecewisesupport started in #433.greater,less,greater_equal,less_equal,cond,logical_or,logical_and,reluandclampusePiecewise, which had no PyTorch or JAX mapping, so exporting any of them raised aKeyError.Piecewisenow maps to nestedtorch.where/jnp.where, with NaN where no branch applies, matching SymPy.LessThan(<=) andGreaterThan(>=) to the strict comparisons, and had no mapping for the strict relations.Max/Minto the reductionsjnp.max/jnp.min, which read the second argument as an axis, somax(x0, x1)raised aTypeError. They now map tojnp.maximum/jnp.minimum.Max/Minwith three or more arguments, which SymPy produces frommax(max(x0, x1), x2), failed on both backends. They are now folded pairwise.extra_torch_mappings, so users could not override them.torch_initialized/jax_initializedwere never set, so setup reran on every export.Tests
test_piecewise_comparison_and_extremum_operators, for both PyTorch and JAX, checks 11 operators against SymPy's NumPy output, on data with ties so<and<=can be told apart.test_extra_mappings_take_precedence.