SymInt and FakeTensor tracing compatibility

Differential Revision: D105054082

Author: ruochen99

generative_recommenders/common.py

@@ -261,9 +261,10 @@ def apply_sampling(
 
 def switch_to_contiguous_if_needed(x: torch.Tensor) -> torch.Tensor:
     if not torch.jit.is_scripting() and torch.compiler.is_compiling():
-        # Tell Dynamo this data-dependent value is in the range (0, 10**9)
-        torch._check(x.size(0) > 0)
-        torch._check(x.size(0) < 10**9)
+        # Range the size as [0, 10**9-1] without requiring s > 0 to be
+        # provable, so unbacked SymInts (e.g. from custom-op fake kernels)
+        # don't fail the strict ``_check(s > 0)`` form at trace time.
+        torch._check_is_size(x.size(0), max=10**9 - 1)
     if x.stride(-1) == 1:
         return x
     return x.contiguous()

generative_recommenders/ops/triton/triton_hstu_linear.py

@@ -32,6 +32,7 @@
 )
 from generative_recommenders.ops.triton.triton_addmm import maybe_triton_addmm_fwd
 from generative_recommenders.ops.utils import maybe_register_custom_op
+from torch.fx.experimental.symbolic_shapes import guard_or_false
 
 
 def _get_layer_norm_mul_dropout_fwd_multirow_configs() -> List[triton.Config]:
@@ -1217,7 +1218,7 @@ def triton_layer_norm_mul_dropout_fwd(
     assert weight.numel() == D
     assert bias.numel() == D
 
-    if N == 0:
+    if guard_or_false(N == 0):
         D = x.shape[1]
         if concat_u and concat_x:
             y = torch.empty((0, 3 * D), dtype=x.dtype, device=x.device)