Repository: jax-ml/jax
Pull request #28526 (2025-05-05 21:02:11)

Replace C++ safe_zip with Python implementation using 3.10 strict keyword
Replace C++ safe_zip with Python implementation using 3.10 strict keyword



Patch:
diff --git a/jax/_src/util.py b/jax/_src/util.py
index 1227b15c5..94b6380b1 100644
--- a/jax/_src/util.py
+++ b/jax/_src/util.py
@@ -42,29 +42,32 @@ T1 = TypeVar("T1")
 T2 = TypeVar("T2")
 T3 = TypeVar("T3")
 
+# safe_zip cannot yet be fully annotated, so we use a strategy similar
+# to that used for builtins.zip in python/typeshed. This supports
+# return types matching input types for up to three arguments.
+@overload
+def safe_zip(__arg1: Iterable[T1], /) -> list[tuple[T1]]: ...
+@overload
+def safe_zip(__arg1: Iterable[T1], __arg2: Iterable[T2], /) -> list[tuple[T1, T2]]: ...
+@overload
+def safe_zip(__arg1: Iterable[T1], __arg2: Iterable[T2], __arg3: Iterable[T3], /) -> list[tuple[T1, T2, T3]]: ...
+@overload
+def safe_zip(__arg1: Iterable[Any], __arg2: Iterable[Any], __arg3: Iterable[Any], __arg4: Iterable[Any], /, *args) -> list[tuple[Any, ...]]: ...
+
+def safe_zip(*args):
+  """
+  Like builtin :func:`zip`, but with additional safety checks.
 
-if TYPE_CHECKING:
-  # safe_zip cannot yet be fully annotated, so we use a strategy similar
-  # to that used for builtins.zip in python/typeshed. This supports
-  # return types matching input types for up to three arguments.
-  @overload
-  def safe_zip(__arg1: Iterable[T1]) -> list[tuple[T1]]: ...
-  @overload
-  def safe_zip(__arg1: Iterable[T1], __arg2: Iterable[T2]) -> list[tuple[T1, T2]]: ...
-  @overload
-  def safe_zip(__arg1: Iterable[T1], __arg2: Iterable[T2], __arg3: Iterable[T3]) -> list[tuple[T1, T2, T3]]: ...
-  @overload
-  def safe_zip(__arg1: Iterable[Any], __arg2: Iterable[Any], __arg3: Iterable[Any], __arg4: Iterable[Any], *args) -> list[tuple[Any, ...]]: ...
+  The differences from :func:`zip` are:
 
-  def safe_zip(*args):
-    args = list(map(list, args))
-    n = len(args[0])
-    for arg in args[1:]:
-      assert len(arg) == n, f'length mismatch: {list(map(len, args))}'
-    return list(zip(*args))
-
-else:
-  safe_zip = jaxlib_utils.safe_zip
+  - :func:`safe_zip` checks that at least one argument is provided.
+  - :func:`safe_zip` checks that all arguments have the same length.
+  - :func:`safe_zip` returns an eagerly-evaluated list instead of a
+    lazily-evaluated iterator.
+  """
+  if not args:
+    raise TypeError("safe_zip requires at least 1 argument.")
+  return list(zip(*args, strict=True))
 
 
 if TYPE_CHECKING:
diff --git a/jaxlib/utils.cc b/jaxlib/utils.cc
index e5bb45e99..1cf679801 100644
--- a/jaxlib/utils.cc
+++ b/jaxlib/utils.cc
@@ -204,98 +204,6 @@ PyMethodDef foreach_def = {
     "ignoring the return values and returns None. The iterables must all have "
     "the same lengths."};
 
-// A variant of zip(...) that:
-// a) returns a list instead of an iterator, and
-// b) checks that the input iterables are of equal length.
-// TODO(phawkins): consider replacing this function with
-// list(zip(..., strict=True)) once TensorFlow 2.13 is released, which should
-// resolve an incompatibility with strict=True and jax2tf.
-PyObject* SafeZip(PyObject* self, PyObject* const* args, Py_ssize_t nargs) {
-  if (nargs < 1) {
-    PyErr_SetString(PyExc_TypeError, "safe_zip requires at least 1 argument");
-    return nullptr;
-  }
-  absl::InlinedVector<nb::object, 4> iterators;
-  iterators.reserve(nargs);
-  for (Py_ssize_t i = 0; i < nargs; ++i) {
-    PyObject* it = PyObject_GetIter(args[i]);
-    if (!it) return nullptr;
-    iterators.push_back(nb::steal<nb::object>(it));
-  }
-
-  // Try to use a length hint to estimate how large a list to allocate.
-  Py_ssize_t length_hint = PyObject_LengthHint(args[0], 2);
-  if (PyErr_Occurred()) {
-    PyErr_Clear();
-  }
-  if (length_hint < 0) {
-    length_hint = 2;
-  }
-
-  nb::list list = nb::steal<nb::list>(PyList_New(length_hint));
-  int n = 0;  // Current true size of the list
-
-  while (true) {
-    nb::object tuple;
-    nb::object v = nb::steal<nb::object>(PyIter_Next(iterators[0].ptr()));
-    if (PyErr_Occurred()) return nullptr;
-
-    if (v.ptr()) {
-      tuple = nb::steal<nb::object>(PyTuple_New(nargs));
-      if (!tuple.ptr()) return nullptr;
-
-      PyTuple_SET_ITEM(tuple.ptr(), 0, v.release().ptr());
-      for (size_t i = 1; i < iterators.size(); ++i) {
-        v = nb::steal<nb::object>(PyIter_Next(iterators[i].ptr()));
-        if (PyErr_Occurred()) return nullptr;
-        if (!v.ptr()) {
-          PyErr_Format(PyExc_ValueError,
-                       "safe_zip() argument %u is shorter than argument 1",
-                       i + 1);
-          return nullptr;
-        }
-        PyTuple_SET_ITEM(tuple.ptr(), i, v.release().ptr());
-      }
-    } else {
-      // No more elements should be left. Checks the other iterators are
-      // exhausted.
-      for (size_t i = 1; i < iterators.size(); ++i) {
-        v = nb::steal<nb::object>(PyIter_Next(iterators[i].ptr()));
-        if (PyErr_Occurred()) return nullptr;
-        if (v.ptr()) {
-          PyErr_Format(PyExc_ValueError,
-                       "safe_zip() argument %u is longer than argument 1",
-                       i + 1);
-          return nullptr;
-        }
-      }
-
-      // If the length hint was too large, truncate the list to the true size.
-      if (n < length_hint) {
-        if (PyList_SetSlice(list.ptr(), n, length_hint, nullptr) < 0) {
-          return nullptr;
-        }
-      }
-      return list.release().ptr();
-    }
-
-    if (n < length_hint) {
-      PyList_SET_ITEM(list.ptr(), n, tuple.release().ptr());
-    } else {
-      if (PyList_Append(list.ptr(), tuple.ptr()) < 0) {
-        return nullptr;
-      }
-      tuple = nb::object();
-    }
-    ++n;
-  }
-}
-
-PyMethodDef safe_zip_def = {
-    "safe_zip",
-    reinterpret_cast<PyCFunction>(SafeZip),
-    METH_FASTCALL,
-};
 
 nb::list TopologicalSort(nb::str parents_attr,
                          nb::iterable end_nodes_iterable) {
@@ -368,8 +276,6 @@ NB_MODULE(utils, m) {
       PyCFunction_NewEx(&safe_map_def, /*self=*/nullptr, module_name.ptr()));
   m.attr("foreach") = nb::steal<nb::object>(
       PyCFunction_NewEx(&foreach_def, /*self=*/nullptr, module_name.ptr()));
-  m.attr("safe_zip") = nb::steal<nb::object>(
-      PyCFunction_NewEx(&safe_zip_def, /*self=*/nullptr, module_name.ptr()));
 
   m.def("topological_sort", &TopologicalSort, nb::arg("parents_attr"),
         nb::arg("end_nodes"),


Test patch:
diff --git a/tests/util_test.py b/tests/util_test.py
index 53414dae9..90506117a 100644
--- a/tests/util_test.py
+++ b/tests/util_test.py
@@ -175,28 +175,28 @@ class SafeZipTest(jtu.JaxTestCase):
     )
 
   def test_safe_zip_errors(self):
-    with self.assertRaisesRegex(
-        TypeError, "safe_zip requires at least 1 argument"
+    with self.assertRaisesWithLiteralMatch(
+        TypeError, "safe_zip requires at least 1 argument."
     ):
       util.safe_zip()
 
-    with self.assertRaisesRegex(
+    with self.assertRaisesWithLiteralMatch(
         TypeError, "'function' object is not iterable"
     ):
       util.safe_zip(lambda x: x)
 
-    with self.assertRaisesRegex(
-        ValueError, r"safe_zip\(\) argument 2 is longer than argument 1"
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, "zip() argument 2 is longer than argument 1"
     ):
       util.safe_zip(range(3), range(4))
 
-    with self.assertRaisesRegex(
-        ValueError, r"safe_zip\(\) argument 2 is shorter than argument 1"
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, "zip() argument 2 is shorter than argument 1"
     ):
       util.safe_zip(range(7), range(2))
 
-    with self.assertRaisesRegex(
-        ValueError, r"safe_zip\(\) argument 2 is longer than argument 1"
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, "zip() argument 2 is longer than argument 1"
     ):
       util.safe_zip((), range(3))