`torch.cond` with tensors created in `[true/false]_fn` fails during `run_decompositions`

### 🐛 Describe the bug

```python
import torch

def branch(x):
torch.tensor(0) # works without this line
return x.clone()

class Module(torch.nn.Module):
def forward(self, x):
return torch.cond(x.any(), branch, branch, (x,))

ep = torch.export.export(Module(), (torch.empty(()),))
ep.run_decompositions()
```

### Error logs

```python
Traceback (most recent call last):
File "bug.py", line 12, in
ep.run_decompositions()
~~~~~~~~~~~~~~~~~~~~~^^
File "/.../lib/python3.14/site-packages/torch/export/exported_program.py", line 124, in wrapper
return fn(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/export/exported_program.py", line 1530, in run_decompositions
return _decompose_exported_program(
self,
...<3 lines>...
decompose_custom_triton_ops=decompose_custom_triton_ops,
)
File "/.../lib/python3.14/site-packages/torch/export/exported_program.py", line 1005, in _decompose_exported_program
) = _decompose_and_get_gm_with_new_signature_constants(
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
ep,
^^^
...<3 lines>...
decompose_custom_triton_ops=decompose_custom_triton_ops,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/.../lib/python3.14/site-packages/torch/export/exported_program.py", line 483, in _decompose_and_get_gm_with_new_signature_constants
aten_export_artifact = _export_to_aten_ir(
patched_mod,
...<6 lines>...
decompose_custom_triton_ops=decompose_custom_triton_ops,
)
File "/.../lib/python3.14/site-packages/torch/export/_trace.py", line 1042, in _export_to_aten_ir
gm, graph_signature = transform(_aot_export_joint_with_descriptors)(
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
stack,
^^^^^^
...<5 lines>...
_record_nn_module_stack=True,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/.../lib/python3.14/site-packages/torch/export/_trace.py", line 971, in _aot_export_joint_with_descriptors
joint_with_descriptors = aot_export_joint_with_descriptors(
stack,
...<4 lines>...
_record_nn_module_stack=_record_nn_module_stack,
)
File "/.../lib/python3.14/site-packages/torch/_functorch/aot_autograd.py", line 1414, in aot_export_joint_with_descriptors
aot_state = create_aot_state(
stack,
...<5 lines>...
shape_env,
)
File "/.../lib/python3.14/site-packages/torch/_functorch/aot_autograd.py", line 582, in create_aot_state
fw_metadata = run_functionalized_fw_and_collect_metadata(
...<4 lines>...
pre_dispatch=aot_config.pre_dispatch,
)(*_dup_fake_script_obj(fake_flat_args))
File "/.../lib/python3.14/site-packages/torch/_functorch/_aot_autograd/collect_metadata_analysis.py", line 220, in inner
flat_f_outs = f(*flat_f_args)
File "/.../lib/python3.14/site-packages/torch/_functorch/_aot_autograd/utils.py", line 192, in flat_fn
tree_out = fn(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/_functorch/_aot_autograd/graph_capture_wrappers.py", line 1534, in functional_call
out = PropagateUnbackedSymInts(mod).run(*args)
File "/.../lib/python3.14/site-packages/torch/fx/interpreter.py", line 197, in run
self.env[node] = self.run_node(node)
~~~~~~~~~~~~~^^^^^^
File "/.../lib/python3.14/site-packages/torch/fx/experimental/symbolic_shapes.py", line 8537, in run_node
result = super().run_node(n)
File "/.../lib/python3.14/site-packages/torch/fx/interpreter.py", line 294, in run_node
return getattr(self, n.op)(n.target, args, kwargs)
~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/fx/interpreter.py", line 377, in call_function
return target(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/_higher_order_ops/cond.py", line 54, in __call__
return super().__call__(pred, true_fn, false_fn, operands)
~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/_ops.py", line 539, in __call__
return self.dispatch(dispatch_key_set.highestPriorityTypeId(), *args, **kwargs)
~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/_ops.py", line 386, in dispatch
return kernel(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/_ops.py", line 336, in maybe_run_autograd
return self(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/_higher_order_ops/cond.py", line 54, in __call__
return super().__call__(pred, true_fn, false_fn, operands)
~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/_ops.py", line 539, in __call__
return self.dispatch(dispatch_key_set.highestPriorityTypeId(), *args, **kwargs)
~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/_ops.py", line 422, in dispatch
result = handler(mode, *args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/_ops.py", line 193, in functionalize_dispatch_mode_fn
return fn(PythonFunctionalizeAPI(mode), *args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/_higher_order_ops/cond.py", line 721, in cond_func
hop_instance = HopInstance.create(cond_op, pred, true_fn, false_fn, inputs)
File "/.../lib/python3.14/site-packages/torch/_higher_order_ops/utils.py", line 1198, in create
return HopInstance(hop, hop.gen_schema(*args, **kwargs))
~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/_higher_order_ops/cond.py", line 61, in gen_schema
then_gm: torch.fx.GraphModule = materialize_as_graph(true_fn, operands)
~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/_higher_order_ops/utils.py", line 1283, in materialize_as_graph
gm = _materialize_as_graph_inner()
File "/.../lib/python3.14/site-packages/torch/_dynamo/eval_frame.py", line 1280, in _fn
return fn(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/_higher_order_ops/utils.py", line 1279, in _materialize_as_graph_inner
return _maybe_reenter_make_fx(
~~~~~~~~~~~~~~~~~~~~~~~
fn, subgraph_decomp_table=subgraph_decomp_table
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
)(*unfunc_t)
~^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/_higher_order_ops/utils.py", line 152, in wrapped
return make_fx(fn, decomposition_table=subgraph_decomp_table)(*args)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^
File "/.../lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 2965, in wrapped
return make_fx_tracer.trace(f, *args)
~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 2867, in trace
return self._trace_inner(f, *args)
~~~~~~~~~~~~~~~~~^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 2828, in _trace_inner
t = dispatch_trace(
wrap_key(func, args, self.fx_tracer, self.pre_dispatch),
tracer=self.fx_tracer,
concrete_args=tuple(phs),
)
File "/.../lib/python3.14/site-packages/torch/_compile.py", line 54, in inner
return disable_fn(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/_dynamo/eval_frame.py", line 1280, in _fn
return fn(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 1673, in dispatch_trace
graph = tracer.trace(root, concrete_args) # type: ignore[arg-type]
File "/.../lib/python3.14/site-packages/torch/_dynamo/eval_frame.py", line 1280, in _fn
return fn(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/fx/_symbolic_trace.py", line 912, in trace
(self.create_arg(fn(*args)),),
~~^^^^^^^
File "/.../lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 1743, in wrapped
out = f(*tensors) # type:ignore[call-arg]
File "", line 1, in
File "/.../lib/python3.14/site-packages/torch/fx/graph_module.py", line 949, in call_wrapped
return self._wrapped_call(self, *args, **kwargs)
~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/fx/graph_module.py", line 461, in __call__
raise e
File "/.../lib/python3.14/site-packages/torch/fx/graph_module.py", line 447, in __call__
return super(self.cls, obj).__call__(*args, **kwargs) # type: ignore[misc]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/fx/_symbolic_trace.py", line 886, in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 1427, in call_module
return forward(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/fx/_symbolic_trace.py", line 879, in forward
return _orig_module_call(mod, *args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
return forward_call(*args, **kwargs)
File ".54 from /.../lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py:1720 in wrapped", line 6, in forward
lift_fresh_copy = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0); _tensor_constant0 = None
File "/.../lib/python3.14/site-packages/torch/_ops.py", line 871, in __call__
return self._op(*args, **kwargs)
~~~~~~~~^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 1798, in __torch_function__
return func(*args, **kwargs)
File "/.../lib/python3.14/site-packages/torch/_ops.py", line 871, in __call__
return self._op(*args, **kwargs)
~~~~~~~~^^^^^^^^^^^^^^^^^
File "/.../lib/python3.14/site-packages/torch/_subclasses/functional_tensor.py", line 280, in __torch_dispatch__
raise RuntimeError(
"Attempting to use FunctionalTensor on its own. Instead, please use it with a corresponding FunctionalTensorMode()"
)
RuntimeError: Attempting to use FunctionalTensor on its own. Instead, please use it with a corresponding FunctionalTensorMode()

While executing %cond : [num_users=1] = call_function[target=torch.ops.higher_order.cond](args = (%any_1, %true_graph_0, %false_graph_0, (%x,)), kwargs = {})
Original traceback:
File "bug.py", line 9, in forward
return torch.cond(x.any(), branch, branch, (x,))
File ".5", line 9, in forward
cond = torch.ops.higher_order.cond(l_args_0_, cond_true_0, cond_false_0, (l_args_3_0_,)); l_args_0_ = cond_true_0 = cond_false_0 = l_args_3_0_ = None
Use tlparse to see full graph. (https://github.com/pytorch/tlparse?tab=readme-ov-file#tlparse-parse-structured-pt2-logs)
```

### Versions

```
Collecting environment information...
PyTorch version: 2.12.0.dev20260414
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 26.3.1 (arm64)
GCC version: Could not collect
Clang version: 17.0.0 (clang-1700.6.3.2)
CMake version: version 4.2.3
Libc version: N/A

Python version: 3.14.2 (v3.14.2:df793163d58, Dec 5 2025, 12:18:06) [Clang 16.0.0 (clang-1600.0.26.6)] (64-bit runtime)
Python platform: macOS-26.3.1-arm64-arm-64bit-Mach-O
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Apple M4 Pro

Versions of relevant libraries:
[pip3] Could not collect
[conda] Could not collect
```

cc @chauhang @penguinwu @avikchaudhuri @zhxchen17 @tugsbayasgalan @angelayi @ydwu4