Revisit as_jax_op#1613

Closed

aseyboldt wants to merge 16 commits intopymc-devs:mainfrom

aseyboldt:as-jax-opt

Commits on Sep 9, 2025

first commit for a decorator that transforms JAX to pytensor

jdehning
authored and
aseyboldt
committed
Add more tests

jdehning
authored and
aseyboldt
committed
Define JAXOp outside of the decorator

jdehning
authored and
aseyboldt
committed
Added comment regarding flattening of inputs

jdehning
authored and
aseyboldt
committed
Add as_jax_op to pytensor.__init__.py and to documentation

jdehning
authored and
aseyboldt
committed
Add [jax] requirement to readthedocs in order to read the docstring of as_jax_op

jdehning
authored and
aseyboldt
committed
Added an example to the docstring of as_jax_op

jdehning
authored and
aseyboldt
committed
Use infer_static_shape, currently still with the possibility to use the previous approach for testing purposes

jdehning
authored and
aseyboldt
committed
Remove `sol` in variable names

jdehning
authored and
aseyboldt
committed
Rename tests and make static variables test more meaningfull

jdehning
authored and
aseyboldt
committed
More test renaming, forgot a few

jdehning
authored and
aseyboldt
committed
Refactoring of ops.py: code is in general cleaner, and JAXOp can now be used without the decorator @as_jax_op

jdehning
authored and
aseyboldt
committed
Clean up tests

jdehning
authored and
aseyboldt
committed
Add to some tests a direct call to JAXOp

jdehning
authored and
aseyboldt
committed
temporary as_jax_op fix
aseyboldt
committed

Commits on Sep 16, 2025

Simplify as_jax_op
aseyboldt
committed