Support symbolic values in torch.arange and other tensor operations

[IvanYashchuk]

🚀 Feature

Thunder needs to support using symbolic scalar values as arguments to tensor creation operations, particularly torch.arange with symbolic start/end bounds.

Current Behavior

Likely requires concrete values for tensor operations.

Expected Behavior

Support symbolic values as arguments to:

• torch.arange(symbolic_start, symbolic_end, ...)
• torch.zeros(symbolic_size, ...)
• .view(symbolic_shape, ...) and other reshaping ops
• .index_copy_(dim, symbolic_indices, ...)

Example from torch.compile

# cache_position: "i64[1]" = torch.arange(cumulative_length, add, device=...)
cache_position: "i64[s50]" = torch.arange(
    l_kwargs_past_key_values_layers_0_cumulative_length, 
    add,  # add is Sym(s50 + s67)
    device=device(type='cuda', index=0)
)

The output tensor has symbolic shape "i64[s50]".

Minimal Reproduction Case

import torch
import thunder

@thunder.jit
def create_cache_position(cumulative_length: int, seq_len: int) -> torch.Tensor:
    """Create cache position indices using symbolic bounds."""
    new_cumulative_length = cumulative_length + seq_len
    
    # torch.arange with symbolic start and end
    cache_position = torch.arange(
        cumulative_length,
        new_cumulative_length,
        device='cuda'
    )
    
    return cache_position

result = create_cache_position(1024, 32)
assert result.shape == (32,)
assert result[0].item() == 1024
assert result[-1].item() == 1055

Related Operations

• torch.arange(symbolic, symbolic)
• torch.zeros(symbolic, *shape)
• tensor.view(-1, symbolic)
• tensor.scatter(1, symbolic_indices, value)

Success Criteria

• torch.arange accepts symbolic start/end
• Output tensor shape is correctly inferred as symbolic
• Other tensor ops support symbolic size arguments
• Generated code handles runtime dimensions correctly

[shino16]

With #2779 and the following patch on nvFuser, torch.arange will work under nvFuser. I'm exploring what other ops that we need to add to nvFuser's concretization process.

diff --git a/csrc/evaluator_common.cpp b/csrc/evaluator_common.cpp
index 1b25d24c6..5e983777b 100644
--- a/csrc/evaluator_common.cpp
+++ b/csrc/evaluator_common.cpp
@@ -604,6 +604,9 @@ void NaiveValueMachine::runUnaryOp(int index) {
     case UnaryOpType::Abs:
       dest = abs(src);
       break;
+    case UnaryOpType::Ceil:
+      dest = ceil(src);
+      break;
     case UnaryOpType::LogicalNot:
       dest = !src;
       break;
diff --git a/csrc/ir/internal_nodes.cpp b/csrc/ir/internal_nodes.cpp
index 0babb967e..9a0baa329 100644
--- a/csrc/ir/internal_nodes.cpp
+++ b/csrc/ir/internal_nodes.cpp
@@ -545,6 +545,9 @@ std::vector<PolymorphicValue> UnaryOp::evaluate(
     case UnaryOpType::Abs:
       return {abs(in)};
       break;
+    case UnaryOpType::Ceil:
+      return {ceil(in)};
+      break;
     case UnaryOpType::LogicalNot:
       return {!in};
       break;
diff --git a/csrc/polymorphic_value.h b/csrc/polymorphic_value.h
index 47cf1456d..1b1cfa34e 100644
--- a/csrc/polymorphic_value.h
+++ b/csrc/polymorphic_value.h
@@ -398,6 +398,19 @@ inline PolymorphicValue abs(const PolymorphicValue& a) {
   NVF_THROW("PolymorphicValue abs not implemented for ", a.type().name());
 }
 
+inline PolymorphicValue ceil(const PolymorphicValue& a) {
+  if (a.is<int64_t>() || a.is<bool>()) {
+    return a;
+  }
+  if (a.is<double>()) {
+    return PolymorphicValue(std::ceil(a.as<double>()));
+  }
+  if (a.is<at::Tensor>()) {
+    return a.as<at::Tensor>().ceil();
+  }
+  NVF_THROW("PolymorphicValue abs not implemented for ", a.type().name());
+}
+
 inline PolymorphicValue erf(const PolymorphicValue& a) {
   if (a.is<at::Tensor>()) {
     return PolymorphicValue(a.as<at::Tensor>().erf());