fix(rng): make default_generators iterable instead of looping forever

tests/integration/ops/test_rng_dispatch.py

@@ -219,6 +219,23 @@ def test_default_generators_populated(self):
             "default generator state is not philox-shaped (2 x int64)"
         )
 
+    @pytest.mark.anyplatform
+    @pytest.mark.main_ops
+    @pytest.mark.parametrize("attr", ["cuda", "flagos"])
+    def test_default_generators_iterable(self, attr):
+        # Upstream types default_generators as a *tuple*, so callers iterate it,
+        # slice it and list() it. Our shims are list-like proxies with no
+        # __iter__, which sends Python to the legacy protocol: __getitem__(0, 1,
+        # 2, ...) until IndexError. Unbounded __getitem__ therefore made
+        # `for g in default_generators` an infinite loop that allocated a fresh
+        # generator per step -- a hang, not an error, so nothing surfaced it.
+        gens = getattr(torch if attr == "cuda" else torch_fl, attr).default_generators
+        n = len(gens)
+        assert len(list(gens)) == n, "iteration does not stop at device_count"
+        assert len(gens[:2]) == min(2, n), "slicing is not supported"
+        with pytest.raises(IndexError):
+            gens[n]
+
     @pytest.mark.anyplatform
     @pytest.mark.main_ops
     def test_manual_seed_reaches_flagos_module(self):

torch_fl/accelerator/cuda/_cuda_compat.py

@@ -85,10 +85,28 @@ class _CudaDefaultGenerators:
     Indexing yields a real (lazily created) per-device CUDA generator; ``len``
     reports the device count so flag_gems' ``len(default_generators) == 0``
     guard is False and it uses the generator instead of erroring.
+
+    Upstream declares ``default_generators`` as a *tuple*, so callers are
+    entitled to iterate it, slice it or wrap it in ``list()``. Bounds-checking
+    ``__getitem__`` is what makes that safe: with no ``__iter__``, Python falls
+    back to the legacy protocol of calling ``__getitem__(0, 1, 2, ...)`` until
+    IndexError, so an unchecked index turned ``for g in default_generators``
+    into an infinite loop that allocated a fresh CUDA generator per step.
     """
 
+    def __iter__(self):
+        return (self[i] for i in range(len(self)))
+
     def __getitem__(self, idx):
-        return _get_cuda_generator(int(idx))
+        n = len(self)
+        if isinstance(idx, slice):
+            return tuple(self[i] for i in range(*idx.indices(n)))
+        idx = int(idx)
+        if idx < 0:  # negative indices wrap, as on the tuple this replaces
+            idx += n
+        if not 0 <= idx < n:
+            raise IndexError(f"device index {idx} out of range for {n} device(s)")
+        return _get_cuda_generator(idx)
 
     def __len__(self):
         try:

torch_fl/accelerator/metax/_metax_compat.py

@@ -182,10 +182,31 @@ def _get_cuda_generator(idx):
 class _CudaDefaultGenerators:
     """list-like stand-in for ``torch.cuda.default_generators`` (see
     _get_cuda_generator). Indexing yields a per-device CUDA generator; ``len``
-    reports the device count so flag_gems' empty-tuple guard is False."""
+    reports the device count so flag_gems' empty-tuple guard is False.
+
+    Upstream declares ``default_generators`` as a *tuple*, so callers are
+    entitled to iterate it, slice it or wrap it in ``list()``. Bounds-checking
+    ``__getitem__`` is what makes that safe: with no ``__iter__``, Python falls
+    back to the legacy protocol of calling ``__getitem__(0, 1, 2, ...)`` until
+    IndexError, so an unchecked index turned ``for g in default_generators``
+    into an infinite loop that allocated a fresh CUDA generator per step.
+    """
+
+    def __iter__(self):
+        return (self[i] for i in range(len(self)))
 
     def __getitem__(self, idx):
-        return _get_cuda_generator(int(idx))
+        n = len(self)
+        if isinstance(idx, slice):
+            return tuple(self[i] for i in range(*idx.indices(n)))
+        idx = int(idx)
+        if idx < 0:  # negative indices wrap, as on the tuple this replaces
+            idx += n
+        if not 0 <= idx < n:
+            raise IndexError(
+                f"device index {idx} out of range for {n} flagos device(s)"
+            )
+        return _get_cuda_generator(idx)
 
     def __len__(self):
         try:

torch_fl/flagos/__init__.py

@@ -146,9 +146,28 @@ def _patched_getitem(self, indices):
 
 # default_generators: list of one Generator per device, required by FlagGems
 class _DefaultGenerators:
-    """Lazy list-like accessor for per-device default generators."""
+    """Lazy list-like accessor for per-device default generators.
+
+    Bounds-checked so iteration terminates: with no ``__iter__``, Python's
+    legacy protocol calls ``__getitem__(0, 1, 2, ...)`` until IndexError, and
+    an out-of-range index otherwise surfaces as a RuntimeError from C++ (which
+    aborts iteration instead of ending it).
+    """
+
+    def __iter__(self):
+        return (self[i] for i in range(len(self)))
 
     def __getitem__(self, device):
+        n = len(self)
+        if isinstance(device, slice):
+            return tuple(self[i] for i in range(*device.indices(n)))
+        device = int(device)
+        if device < 0:  # negative indices wrap, as on a list
+            device += n
+        if not 0 <= device < n:
+            raise IndexError(
+                f"device index {device} out of range for {n} flagos device(s)"
+            )
         return _C._get_default_generator(device)
 
     def __len__(self):