[torchax] Remove safe_zip (#9525)

Author: bhavya01

torchax/torchax/interop.py

@@ -225,8 +225,7 @@ def j2t_autograd(fn, call_jax=call_jax):
 
   @wraps(fn)
   def inner(*args, **kwargs):
-    from jax.tree_util import tree_flatten, tree_unflatten
-    from jax.util import safe_zip
+    from jax.tree_util import tree_flatten
 
     class JaxFun(torch.autograd.Function):
 
@@ -261,8 +260,8 @@ def backward(ctx, *grad_out):
         # The subsequent gradients correspond to flat_inputs.
         # We need to put a None for inputs that did not require gradients.
         final_grads = [None]
-        for needs_grad, grad in safe_zip(ctx.needs_input_grad[1:],
-                                         input_grads_structured):
+        for needs_grad, grad in zip(
+            ctx.needs_input_grad[1:], input_grads_structured, strict=True):
           final_grads.append(grad if needs_grad else None)
 
         return tuple(final_grads)