3 ms·

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 p

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.

sillysaurusx 5y ago
This is entirely possible. You have direct access to XLA, after all: https://jax.readthedocs.io/en/latest/notebooks/XLA_in_Python.html 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#setdimensionsize https://www.tensorflow.org/xla/operation_semantics#setdimens...