The TensorFlow backend currently returns an eager lambdified function: it is neither decorated with tf.function nor compiled with XLA. JAX functions, by contrast, are wrapped with jax.jit by default. This makes the default backend behavior asymmetric and can leave TensorFlow significantly slower for element-wise expressions because every operation is dispatched separately and intermediate tensors are materialized.
The implementation should distinguish between TensorFlow graph tracing (tf.function) and XLA compilation (tf.function(jit_compile=True)). Unconditionally enabling XLA may not be optimal for every expression or workload and can reject operations that XLA does not support. The existing use_jit option should have clear, consistent semantics across backends, with an appropriate default and an escape hatch where needed.
Performance should be evaluated with backend-native inputs and explicit synchronization, separating first-call tracing/compilation cost from warmed execution. Benchmarks should cover both 32-bit and 64-bit precision and avoid including NumPy conversion or device-to-host transfer in the measured operation.
Previous implementation
#441 implemented this by unconditionally wrapping every TensorFlow lambdified function in tf.function(..., jit_compile=True). The PR had no negative review and its tests passed, but its CI benchmark appeared to become slower. The change was therefore left open and eventually closed without merging. The approach was not fundamentally invalid, but it was too unconditional and did not investigate graph mode without XLA, compilation versus steady-state timing, unsupported operations, retracing, or representative workload-dependent performance.
Acceptance criteria
- Define the TensorFlow behavior of
use_jit, including whether it selects graph mode or XLA compilation.
- Add correctness tests for eager and compiled execution where applicable.
- Benchmark eager execution,
tf.function, and XLA separately on representative workloads, excluding tracing/compilation and host conversions from steady-state timings.
- Measure or document first-call tracing and compilation cost separately.
- Verify behavior in both
float32 and float64 modes.
- Preserve a fallback for expressions or platforms that cannot use XLA.
This issue is part of #566. Compiling the innermost TensorFlow function complements the larger effort in #570 to make complete estimators pure and JIT-compilable.
The TensorFlow backend currently returns an eager lambdified function: it is neither decorated with
tf.functionnor compiled with XLA. JAX functions, by contrast, are wrapped withjax.jitby default. This makes the default backend behavior asymmetric and can leave TensorFlow significantly slower for element-wise expressions because every operation is dispatched separately and intermediate tensors are materialized.The implementation should distinguish between TensorFlow graph tracing (
tf.function) and XLA compilation (tf.function(jit_compile=True)). Unconditionally enabling XLA may not be optimal for every expression or workload and can reject operations that XLA does not support. The existinguse_jitoption should have clear, consistent semantics across backends, with an appropriate default and an escape hatch where needed.Performance should be evaluated with backend-native inputs and explicit synchronization, separating first-call tracing/compilation cost from warmed execution. Benchmarks should cover both 32-bit and 64-bit precision and avoid including NumPy conversion or device-to-host transfer in the measured operation.
Previous implementation
#441 implemented this by unconditionally wrapping every TensorFlow lambdified function in
tf.function(..., jit_compile=True). The PR had no negative review and its tests passed, but its CI benchmark appeared to become slower. The change was therefore left open and eventually closed without merging. The approach was not fundamentally invalid, but it was too unconditional and did not investigate graph mode without XLA, compilation versus steady-state timing, unsupported operations, retracing, or representative workload-dependent performance.Acceptance criteria
use_jit, including whether it selects graph mode or XLA compilation.tf.function, and XLA separately on representative workloads, excluding tracing/compilation and host conversions from steady-state timings.float32andfloat64modes.This issue is part of #566. Compiling the innermost TensorFlow function complements the larger effort in #570 to make complete estimators pure and JIT-compilable.