[RL][Flex] batch invariant Flex Attention doesn't work with general compile

titan is missing the mechanism in compile that uses aot_eager backend overall but lets you scoop a region to plug in inductor backend supported code (ie Flex attention) **additional context** if i use `compile=CompileConfig(enable=True, backend = aot_eager)` in `/home/liangel/local/torchtitan/torchtitan/experiments/rl/config_registry.py` `rl_grpo_qwen3_0_6b_flex_batch_invariant()` config, i don't get matching numerics between trainer and generator (`bit_wise/logprob_diff/max: 0.44`). ``` Train | Step: 1 loss/mean: 0.00002 bit_wise/logprob_diff/max: 0.44 reward/_mean: 0.51 reward/_max: 1.30 reward/zero_std_frac: 0 rollout/response_length/max: 100.0 train/grad_norm/mean: 0.73 train/lr: 1.0e-06 perf/tokens_per_second: 1238.6 timing/step: 16.15 ``` if i use inductor backend `compile=CompileConfig(enable=True)` i get ``` File "/home/liangel/local/pytorch/torch/_dynamo/eval_frame.py", line 1143, in compile_wrapper result = fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/home/liangel/local/pytorch/torch/nn/modules/module.py", line 1789, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/liangel/local/torchtitan/torchtitan/models/qwen3/model.py", line 54, in forward def forward( File "/home/liangel/local/pytorch/torch/_dynamo/eval_frame.py", line 1421, in _fn return fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/home/liangel/local/pytorch/torch/_functorch/aot_autograd.py", line 1277, in forward return compiled_fn(full_args) ^^^^^^^^^^^^^^^^^^^^^^ File "/home/liangel/local/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 2549, in __call__ return self.compiled_fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/liangel/local/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1142, in runtime_wrapper result = _codegen_runtime_wrapper( ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/liangel/local/pytorch/torch/_functorch/_aot_autograd/subclass_codegen.py:codegen(runtime_wrapper_orchestration)", line 8, in _runtime_wrapper File "/home/liangel/local/pytorch/torch/_functorch/_aot_autograd/subclass_codegen.py:codegen(subclass_wrapper)", line 11, in inner_fn AttributeError: 'Tensor' object has no attribute 'elem' ``` cc @tianyu-l @SherlockNoMad