Support subsequence that does not support slice as index.

ml_metrics/_src/utils/iter_utils.py

@@ -157,11 +157,11 @@ def slice(self, slice_: slice) -> Iterator[_ValueT]:
     if start.seq_idx == stop.seq_idx:
       return itt.islice(self._sequences[start.seq_idx], start.idx, stop.idx)
     # Chain multiple sequences together with correct slices.
-    sequences = [self._sequences[start.seq_idx][start.idx :]]
+    sequences = [itt.islice(self._sequences[start.seq_idx], start.idx, None)]
     for i_seq in range(start.seq_idx + 1, stop.seq_idx):
       sequences.append(self._sequences[i_seq])
     if stop.idx:
-      sequences.append(self._sequences[stop.seq_idx][: stop.idx])
+      sequences.append(itt.islice(self._sequences[stop.seq_idx], stop.idx))
     return itt.chain.from_iterable(sequences)
 
   def __getitem__(self, index) -> _ValueT | Iterator[_ValueT]:

ml_metrics/_src/utils/iter_utils_test.py

@@ -60,6 +60,18 @@ def __iter__(self):
     return iter(self._iteratable)
 
 
+class MockSequence:
+
+  def __init__(self, data):
+    self._data = data
+
+  def __len__(self):
+    return len(self._data)
+
+  def __getitem__(self, i):
+    return self._data[i.__index__()]
+
+
 def range_with_return(n, sleep: float = 0):
   for i in range(n):
     if sleep:
@@ -146,8 +158,18 @@ def test_sequence_array_index_raises(self):
     with self.assertRaisesRegex(ValueError, 'not in array'):
       a.index(10)
 
-  def test_merged_sequences_default(self):
-    a = iter_utils.MergedSequences([range(10), range(10, 20)])
+  @parameterized.named_parameters([
+      dict(
+          testcase_name='default',
+          seqs=[range(10), range(10, 20)],
+      ),
+      dict(
+          testcase_name='non_sliceable_sequence',
+          seqs=[MockSequence(range(10)), MockSequence(range(10, 20))],
+      ),
+  ])
+  def test_merged_sequences_default(self, seqs):
+    a = iter_utils.MergedSequences(seqs)
     self.assertLen(a, 20)
     self.assertEqual(list(a), list(range(20)))
     self.assertEqual(11, a[11])