Merge pull request #21146 from jakevdp:fix-multidot

PiperOrigin-RevId: 632142647
google · May 9, 2024 · 2be3f6d · 2be3f6d
2 parents 89d25bb + 5edfaa6
commit 2be3f6d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
@@ -2017,4 +2017,4 @@ def multi_dot(arrays: Sequence[ArrayLike], *, precision: PrecisionLike = None) -
   if arrs[-1].ndim == 1:
     einsum_axes[-1] = einsum_axes[-1][:1]
   return jnp.einsum(*itertools.chain(*zip(arrs, einsum_axes)),  # type: ignore[arg-type, call-overload]
-                    optimize='optimal', precision=precision)
+                    optimize='auto', precision=precision)