Commit b3b6861
committed
Only MIT-MOT require working on buffers directly
Using JAX Scan machinery to create MIT-SOT, SIT-SOT, and NIT-SOT buffers for us seems to be more performant than working directly on the pre-allocated buffers and reading/writing at every iteration. There is no machinery to work with MIT-MOT directly (just like in PyTensor user-facing Scan).1 parent 14e6c78 commit b3b6861
File tree
2 files changed
+207
-256
lines changed- pytensor/link/jax/dispatch
- tests/link/jax
2 files changed
+207
-256
lines changed
0 commit comments