Use optimize='auto' for multi_dot.

fantasycz · Nov 22, 2024 · 236d4c6 · 236d4c6
1 parent 34a2f0c
commit 236d4c6
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