code: def _safe_zip(*args):
      length = len(args[0])
      assert all(len(a) == length for a in args[1:])
      return zip(*args)
  
file_path: equinox/tests/test_finalise_jaxpr.py
imports: from typing import cast
  import equinox.internal
  import jax
  import jax.core
  import jax.lax
  import jax.numpy
  from .helpers import tree_allclose
label_desc: Use image node_id 1 for calling a global function with example usage: _safe_zip() and returns: zip