Revert "Switch to using Python nested int (pytorch#141166)"

This reverts commit e2e8a7f.

Reverted pytorch#141166 on behalf of https://github.com/clee2000 due to broke docs [GH job link](https://github.com/pytorch/pytorch/actions/runs/11980936976/job/33406870951) [HUD commit link](https://hud.pytorch.org/pytorch/pytorch/commit/e2e8a7fa2e519433a4ec1071f80d2f6f843c6300) ([comment](pytorch#141166 (comment)))

test/test_nestedtensor.py

@@ -8574,101 +8574,6 @@ def f(*args, **kwargs):
                     self.assertEqualNoncontigAware(grads_compile, grads_ref)
 
 
-from torch.nested._internal.nested_int import NestedIntNode
-
-
-class TestNestedInt(torch.testing._internal.common_utils.TestCase):
-    def test_comparisons(self):
-        a = torch.SymInt(NestedIntNode(1, 1))
-        b = torch.SymInt(NestedIntNode(1, 1))
-        c = torch.SymInt(NestedIntNode(2, 1))
-        d = 3
-
-        self.assertTrue(a == a)
-        self.assertTrue(a == b)
-        self.assertFalse(a != a)
-        self.assertFalse(a != b)
-        self.assertFalse(a == c)
-        self.assertTrue(a != c)
-
-        self.assertFalse(a == d)
-        self.assertTrue(a != d)
-        self.assertFalse(d == a)
-        self.assertTrue(d != a)
-
-        # ge
-        self.assertTrue(a >= a)
-        self.assertTrue(a >= b)
-        self.assertTrue(b >= a)
-        with self.assertRaises(ValueError):
-            _ = a >= c
-        with self.assertRaises(ValueError):
-            _ = c >= a
-        with self.assertRaises(ValueError):
-            _ = c >= 3
-        self.assertTrue(c >= 2)
-        self.assertTrue(c >= 1)
-        self.assertFalse(c <= 1)
-
-        # lt
-        self.assertFalse(a < a)
-        self.assertFalse(a < b)
-        self.assertFalse(b < a)
-        with self.assertRaises(ValueError):
-            _ = a < c
-        with self.assertRaises(ValueError):
-            _ = c < a
-        with self.assertRaises(ValueError):
-            _ = 3 < a
-        with self.assertRaises(ValueError):
-            _ = 2 < a
-        self.assertTrue(a > 1)
-
-        # le
-        self.assertTrue(a <= a)
-        self.assertTrue(b <= a)
-        self.assertTrue(a <= b)
-        with self.assertRaises(ValueError):
-            _ = a <= c
-        with self.assertRaises(ValueError):
-            _ = c <= a
-        with self.assertRaises(ValueError):
-            _ = 3 <= c
-        self.assertTrue(c >= 2)
-        self.assertTrue(c >= 1)
-        self.assertFalse(c <= 1)
-
-        # gt
-        self.assertFalse(a > a)
-        self.assertFalse(b > a)
-        self.assertFalse(a > b)
-        with self.assertRaises(ValueError):
-            _ = a > c
-        with self.assertRaises(ValueError):
-            _ = c > a
-        with self.assertRaises(ValueError):
-            _ = a > 3
-        with self.assertRaises(ValueError):
-            _ = a > 2
-        self.assertTrue(a > 1)
-
-    def test_with_factor(self):
-        a = torch.SymInt(NestedIntNode(1, 5))
-        b = torch.SymInt(NestedIntNode(1, 10))
-        # eq
-        self.assertFalse(a == b)
-        self.assertFalse(a >= b)
-        self.assertTrue(b >= a)
-        self.assertTrue(a <= b)
-        self.assertFalse(b <= a)
-        # ne
-        self.assertTrue(a != b)
-        # mul
-        self.assertTrue(a * 2 == b)
-        self.assertTrue(a * 3 >= b)
-        self.assertTrue(a * 2 == 2 * a)
-
-
 instantiate_parametrized_tests(TestNestedTensor)
 instantiate_device_type_tests(TestNestedTensorDeviceType, globals())
 instantiate_device_type_tests(TestNestedTensorAutograd, globals())

torch/_subclasses/fake_tensor.py

@@ -2513,13 +2513,12 @@ def create_symbolic_nested_int(
         # See Note: [Creating symbolic nested int]
         # Returned nested int always has coeff=1; multiply the result by coeff if needed
         import torch.nested._internal.nested_tensor
-        from torch.nested._internal.nested_int import NestedIntNode
 
         if nt_tensor_id is None:
             nt_tensor_id = self.nt_tensor_id_counter
             assert self.enter_stack, "should only called while FakeTensorMode is active"
             self.nt_tensor_id_counter += 1
-        hint = torch.SymInt(NestedIntNode(nt_tensor_id, 1))
+        hint = torch._C._get_nested_int(nt_tensor_id, 1)
 
         src = torch._dynamo.source.EphemeralSource("intermediate_offsets_or_lengths")
         assert self.shape_env is not None

torch/nested/_internal/nested_tensor.py

@@ -5,7 +5,6 @@
 import torch
 from torch._C import DispatchKey, DispatchKeySet
 from torch._prims_common import is_expandable_to
-from torch.nested._internal.nested_int import NestedIntNode
 from torch.utils.weak import WeakTensorKeyDictionary
 
 
@@ -26,7 +25,7 @@ def get_tensor_symint(tensor, *, coeff=1):
 
     tensor_symint = _tensor_symint_registry.get(tensor)
     if tensor_symint is None:
-        tensor_symint = torch.SymInt(NestedIntNode(_tensor_id_counter, coeff))
+        tensor_symint = torch._C._get_nested_int(_tensor_id_counter, coeff)
         _tensor_id_counter += 1
         _tensor_symint_registry[tensor] = tensor_symint
     return tensor_symint