Re:Useful algorithms that are not optimized by Jax, PyTorch,...
by awaythrowact on Thursday July 22, 2021 @02:27PM
(#27922486)
Makes sense. I don't use TF Eager, but I do use Jax, and Jax lets you arbitrarily compose JITed and non-JITed code, which made me think that might be a viable pattern. I guess I wondered if there might be something like "nonstatic_jit(foo)" that would do "julia style" compiling on function foo, in addition to "jit(foo)" that compiles foo to optimized XLA ops. Probably impractical. Thanks.
[ Parent
| on HN ]
Re:Useful algorithms that are not optimized by Jax, PyTorch,...
by sillysaurusx on Thursday July 22, 2021 @10:20PM
(#27926524)
This is entirely possible. You have direct access to XLA, after all: https://jax.readthedocs.io/en/latest/notebooks/XLA_in_Python...
And XLA has dynamic shape semantics (currently unused by jax) via SetDimensionSize: https://www.tensorflow.org/xla/operation_semantics#setdimens...