SAM (Segment Anything Model) failing to export

🐛 Describe the bug

Exporting the SAM model (https://huggingface.co/facebook/sam-vit-base) using torch.export.export() prior to executorch's to_edge() due to data-dependent expressions in the model graph.

Traceback (most recent call last):
  File "/home/collinallen/projects/executorch_enablement/export_models.py", line 226, in <module>
    export_model(name, hf_id, cls, wrapper)
  File "/home/collinallen/projects/executorch_enablement/export_models.py", line 208, in export_model
    exported_program = torch.export.export(wrapped, inputs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/__init__.py", line 319, in export
    raise e
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/__init__.py", line 286, in export
    return _export(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 1152, in wrapper
    raise e
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 1118, in wrapper
    ep = fn(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/exported_program.py", line 123, in wrapper
    return fn(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 2164, in _export
    ep = _export_for_training(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 1152, in wrapper
    raise e
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 1118, in wrapper
    ep = fn(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/exported_program.py", line 123, in wrapper
    return fn(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 2027, in _export_for_training
    export_artifact = export_func(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 1969, in _non_strict_export
    aten_export_artifact = _to_aten_func(  # type: ignore[operator]
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 1763, in _export_to_aten_ir_make_fx
    gm, graph_signature = transform(_make_fx_helper)(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 1892, in _aot_export_non_strict
    gm, sig = aot_export(wrapped_mod, args, kwargs=kwargs, **flags)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 1678, in _make_fx_helper
    gm = make_fx(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2351, in wrapped
    return make_fx_tracer.trace(f, *args)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2283, in trace
    return self._trace_inner(f, *args)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2254, in _trace_inner
    t = dispatch_trace(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/_compile.py", line 53, in inner
    return disable_fn(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 975, in _fn
    return fn(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1283, in dispatch_trace
    graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1865, in trace
    res = super().trace(root, concrete_args)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 850, in trace
    (self.create_arg(fn(*args)),),
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1341, in wrapped
    out = f(*tensors)  # type:ignore[call-arg]
  File "<string>", line 1, in <lambda>
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 1582, in wrapped_fn
    return tuple(flat_fn(*args))
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 184, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 1138, in functional_call
    out = mod(*args[params_len:], **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 825, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1935, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 542, in call_module
    ret_val = forward(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 818, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/export/_trace.py", line 1876, in forward
    tree_out = mod(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 825, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1935, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 542, in call_module
    ret_val = forward(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 818, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/collinallen/projects/executorch_enablement/export_models.py", line 44, in forward
    return self.model(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 825, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1935, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 542, in call_module
    ret_val = forward(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 818, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/transformers/utils/generic.py", line 943, in wrapper
    output = func(self, *args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/transformers/models/sam/modeling_sam.py", line 1486, in forward
    low_res_masks, iou_predictions, mask_decoder_attentions = self.mask_decoder(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 825, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1935, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 542, in call_module
    ret_val = forward(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 818, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/transformers/models/sam/modeling_sam.py", line 534, in forward
    if sparse_prompt_embeddings.sum().item() != 0:
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/__init__.py", line 763, in __bool__
    return self.node.bool_()
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/sym_node.py", line 616, in bool_
    return self.guard_bool("", 0)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/sym_node.py", line 538, in guard_bool
    r = self.evaluate()
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/sym_node.py", line 512, in evaluate
    return self.shape_env.evaluate_sym_node(self, size_oblivious)
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 7237, in evaluate_sym_node
    return self.evaluate_expr(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 7337, in evaluate_expr
    return self._inner_evaluate_expr(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/recording.py", line 272, in wrapper
    return retlog(fn(*args, **kwargs))
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 7360, in _inner_evaluate_expr
    return self._evaluate_expr(
  File "/home/collinallen/miniconda3/envs/executorch/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 7584, in _evaluate_expr
    raise self._make_data_dependent_error(
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not guard on data-dependent expression Ne(zuf0, 0.0) (unhinted: Ne(zuf0, 0.0)).  (Size-like symbols: none)

Caused by: (transformers/models/sam/modeling_sam.py:534 in forward)
For more information, run with TORCH_LOGS="dynamic"
For extended logs when we create symbols, also add TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="zuf0"
If you suspect the guard was triggered from C++, add TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
For more debugging help, see https://docs.google.com/document/d/1HSuTTVvYH1pTew89Rtpeu84Ht3nQEFTYhAX3Ypa_xJs/edit?usp=sharing

For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1

The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your export() call with draft_export().

This failure comes in the pytorch code prior to executorch, so this may not be the appropriate place. Considering this is also a model provided by Meta, I would like to know if there is guidance for handling models with data dependent expressions in the inference graph.

Versions

Collecting environment information...
PyTorch version: 2.9.0.dev20250714+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.2 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0
Clang version: Could not collect
CMake version: version 3.31.6
Libc version: glibc-2.39

Python version: 3.10.18 (main, Jun 5 2025, 13:14:17) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39
Is CUDA available: False
CUDA runtime version: 12.9.86
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4090
Nvidia driver version: 561.17
cuDNN version: Could not collect
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 46 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: GenuineIntel
Model name: 13th Gen Intel(R) Core(TM) i9-13900
CPU family: 6
Model: 183
Thread(s) per core: 2
Core(s) per socket: 8
Socket(s): 1
Stepping: 1
BogoMIPS: 3993.59
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves umip gfni vaes vpclmulqdq rdpid fsrm md_clear flush_l1d arch_capabilities
Hypervisor vendor: Microsoft
Virtualization type: full
L1d cache: 384 KiB (8 instances)
L1i cache: 256 KiB (8 instances)
L2 cache: 16 MiB (8 instances)
L3 cache: 36 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Unknown: No mitigations
Vulnerability Reg file data sampling: Mitigation; Clear Register File
Vulnerability Retbleed: Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI BHI_DIS_S
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected

Versions of relevant libraries:
[pip3] executorch==0.8.0a0+e19851c
[pip3] numpy==2.2.6
[pip3] nvidia-cublas-cu12==12.6.4.1
[pip3] nvidia-cuda-cupti-cu12==12.6.80
[pip3] nvidia-cuda-nvrtc-cu12==12.6.77
[pip3] nvidia-cuda-runtime-cu12==12.6.77
[pip3] nvidia-cudnn-cu12==9.5.1.17
[pip3] nvidia-cufft-cu12==11.3.0.4
[pip3] nvidia-curand-cu12==10.3.7.77
[pip3] nvidia-cusolver-cu12==11.7.1.2
[pip3] nvidia-cusparse-cu12==12.5.4.2
[pip3] nvidia-cusparselt-cu12==0.6.3
[pip3] nvidia-nccl-cu12==2.26.2
[pip3] nvidia-nvjitlink-cu12==12.6.85
[pip3] nvidia-nvtx-cu12==12.6.77
[pip3] pytorch_tokenizers==0.1.0
[pip3] torch==2.9.0.dev20250714+cpu
[pip3] torchao==0.12.0+gitaee079503
[pip3] torchaudio==2.8.0.dev20250714+cpu
[pip3] torchdata==0.11.0
[pip3] torchsr==1.0.4
[pip3] torchtune==0.6.1
[pip3] torchvision==0.24.0.dev20250714+cpu
[pip3] triton==3.3.1
[conda] executorch 0.8.0a0+e19851c pypi_0 pypi
[conda] numpy 2.2.6 pypi_0 pypi
[conda] nvidia-cublas-cu12 12.6.4.1 pypi_0 pypi
[conda] nvidia-cuda-cupti-cu12 12.6.80 pypi_0 pypi
[conda] nvidia-cuda-nvrtc-cu12 12.6.77 pypi_0 pypi
[conda] nvidia-cuda-runtime-cu12 12.6.77 pypi_0 pypi
[conda] nvidia-cudnn-cu12 9.5.1.17 pypi_0 pypi
[conda] nvidia-cufft-cu12 11.3.0.4 pypi_0 pypi
[conda] nvidia-curand-cu12 10.3.7.77 pypi_0 pypi
[conda] nvidia-cusolver-cu12 11.7.1.2 pypi_0 pypi
[conda] nvidia-cusparse-cu12 12.5.4.2 pypi_0 pypi
[conda] nvidia-cusparselt-cu12 0.6.3 pypi_0 pypi
[conda] nvidia-nccl-cu12 2.26.2 pypi_0 pypi
[conda] nvidia-nvjitlink-cu12 12.6.85 pypi_0 pypi
[conda] nvidia-nvtx-cu12 12.6.77 pypi_0 pypi
[conda] pytorch-tokenizers 0.1.0 pypi_0 pypi
[conda] torch 2.9.0.dev20250714+cpu pypi_0 pypi
[conda] torchao 0.12.0+gitaee079503 pypi_0 pypi
[conda] torchaudio 2.8.0.dev20250714+cpu pypi_0 pypi
[conda] torchdata 0.11.0 pypi_0 pypi
[conda] torchsr 1.0.4 pypi_0 pypi
[conda] torchtune 0.6.1 pypi_0 pypi
[conda] torchvision 0.24.0.dev20250714+cpu pypi_0 pypi
[conda] triton 3.3.1 pypi_0 pypi

cc @JacobSzwejbka @angelayi

Reproduce the failure from export_models.py, especially the wrapper at line 44 and the torch.export.export call at line 208, before to_edge(). Trace the SAM mask decoder path reported at transformers/models/sam/modeling_sam.py line 534, then verify that SAM export completes without the data-dependent-expression error.