Fix kd.index docstring an argument check.

PiperOrigin-RevId: 713682449
Change-Id: If5b3d709dc883faf3c1e49a98e84948465263912

Author: odimka

py/koladata/operators/slices.py

@@ -811,7 +811,8 @@ def index(x, dim=-1):
 
   Args:
     x: A DataSlice.
-    dim: The dimension to compute indices over. Requires abs(dim) < get_ndim(x).
+    dim: The dimension to compute indices over.
+      Requires -get_ndim(x) <= dim < get_ndim(x).
       If dim < 0 then dim = get_ndim(x) + dim.
   """
   x = assertion.with_assertion(
@@ -825,7 +826,7 @@ def index(x, dim=-1):
   ndim = assertion.with_assertion(
       ndim,
       (ndim < get_ndim(x)) & (ndim >= 0),
-      'kde.slices.index: expected 0 <= dim < rank',
+      'kde.slices.index: expected -get_ndim(x) <= dim < get_ndim(x)',
   )
 
   aggregated = masking.agg_has(x, ndim)

py/koladata/operators/tests/slices_index_test.py

@@ -133,7 +133,9 @@ def test_out_of_bounds_ndim_error(self, ndim):
     with self.assertRaisesRegex(
         # TODO: b/375621456 - Raise KodaError.
         ValueError,
-        'kde.slices.index: expected 0 <= dim < rank',
+        re.escape(
+            'kde.slices.index: expected -get_ndim(x) <= dim < get_ndim(x)'
+        ),
     ):
       expr_eval.eval(kde.slices.index(x, ndim))