Fix tf.function retracing in TensorFlow benchmark (#27665)
## Summary
- Move `tf.function`-decorated forward functions out of the inner
benchmark loop to prevent unnecessary graph retracing on every
`(batch_size, sequence_length)` iteration
- Update deprecated `experimental_compile` to `jit_compile` (available
since TF 2.4)
- Hoist `import random` out of the inner loop
Fixes #14953
## Motivation
When `run_with_tf_optimizations` is used as a decorator inside the
innermost `(batch_size, sequence_length)` loop, each iteration creates a
new Python function object. Since `tf.function` keys its trace cache on
function identity, a new object means a forced retrace every iteration —
the cached graph is never reused. This defeats the purpose of
`tf.function` and adds significant overhead from repeated graph
construction and optimization passes.
The [TensorFlow documentation on
tracing](https://www.tensorflow.org/guide/function#rules_of_tracing)
explicitly warns against defining `tf.function`-decorated functions
inside loops.
## Changes
**`onnxruntime/python/tools/transformers/benchmark.py`** (1 file, ~35
insertions / ~31 deletions):
1. **Hoisted forward function definitions** (`encoder_forward`,
`encoder_decoder_forward`, `lxmert_forward`) from the inner `batch_size
× sequence_length` loop to the per-model scope. They are now defined
once per model, and the `@run_with_tf_optimizations` decorator (which
applies `@tf.function`) is only invoked once per model.
2. **Changed forward functions to accept `input_ids` as a parameter**
instead of closing over the loop variable. This lets `tf.function` trace
based on the tensor's `(dtype, shape)` spec and reuse cached concrete
functions when shapes repeat across iterations.
3. **Updated `experimental_compile=use_xla`** to
**`jit_compile=use_xla`**. The `experimental_compile` parameter was
deprecated in TF 2.4 (Dec 2020) and removed in TF 2.12.
4. **Moved `import random`** from the innermost loop body to before the
outer model loop — the module only needs to be imported once.
5. **Moved inference function selection** (`if config.is_encoder_decoder
... elif isinstance(config, LxmertConfig) ...`) outside the
batch/sequence loops since it depends only on the model config, not on
batch size or sequence length. The original priority order
(`is_encoder_decoder` checked before `LxmertConfig`) is preserved.
## Test Plan
- [x] `lintrunner -a` passes cleanly (no RUFF or RUFF-FORMAT violations)
- [x] `python -m py_compile benchmark.py` — syntax verified
- [x] Change is purely structural — function behavior (inputs, outputs,
control flow) is identical
- [ ] Manual verification with TensorFlow installed (TF is an optional
dependency not present in the standard CI matrix; this code path is
exercised via `python benchmark.py -e tensorflow`)
