[Grappler] Add Einsum to auto_mixed_precision whitelist

- Einsum is used in some of the official NLP models, and without
  whitelisting it they see no speedup from auto_mixed_precision.
- The op calls into blas gemm routines, and so it should be treated
  the same as MatMul.
This commit is contained in:
Ben Barsdell 2020-03-26 11:46:32 +11:00
parent fd05051846
commit 0107439506

View File

@ -69,6 +69,7 @@ class AutoMixedPrecisionLists {
"CudnnRNNBackpropV3",
"CudnnRNNV2",
"CudnnRNNV3",
"Einsum",
"GRUBlockCell",
"GRUBlockCellGrad",
"LSTMBlockCell",