Merge pull request #25055 from dfm:multi-dot

PiperOrigin-RevId: 702039013
jax-ml · Dec 2, 2024 · 46c748b · 46c748b
2 parents 385328b + 236d4c6
commit 46c748b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
@@ -2119,7 +2119,7 @@ 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[call-overload]
-                    optimize='optimal', precision=precision)
+                    optimize='auto', precision=precision)
 
 
 @export