Implement `wrap_jax` and rename `as_op` to `wrap_py`#1614

Merged

ricardoV94 merged 30 commits intopymc-devs:mainfrom

aseyboldt:as-jax-opt2

Oct 2, 2025

Commits on Sep 16, 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
Simplify as_jax_op
aseyboldt
committed
optionally eval shapes in as_jax_op
aseyboldt
committed
minor coding style changes in as_jax_op
aseyboldt
committed
set output shape to None if not statically known in as_jax_op
aseyboldt
committed
clean up global import
aseyboldt
committed
remove name from JaxOp.__props__
aseyboldt
committed
don't compile in shape eval of as_jax_op
aseyboldt
committed
more tests for as_jax_op
aseyboldt
committed

Commits on Sep 24, 2025

Commits on Sep 30, 2025