Merge pull request #2786 from google/einsum-optimize-default

try optimize=True default with einsum
jax-ml · Apr 22, 2020 · ec03f8e · ec03f8e
2 parents 6db1f0c + 783e24f
commit ec03f8e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
@@ -2558,7 +2558,7 @@ def tensordot(a, b, axes=2, precision=None):
 
 @_wraps(onp.einsum, lax_description=_PRECISION_DOC)
 def einsum(*operands, **kwargs):
-  optimize = kwargs.pop('optimize', 'auto')
+  optimize = kwargs.pop('optimize', True)
   optimize = 'greedy' if optimize is True else optimize
   precision = kwargs.pop('precision', None)
   if kwargs: