commit 92ddee60c3f4bed5004adfd597adeaf8450669f4
author Benjamin Kramer <kramerb@google.com> Wed May 12 02:44:04 2021 -0700
committer Copybara-Service <iree-copybara-bot@google.com> Wed May 12 02:44:37 2021 -0700
tree d302f4502a08946b406de1cd6521ff57a15d2888
parent 619265c937b494068457d0b802d6486be8512c2e

[TF:XLA:MLIR] Add support for tf.Mean on ranked dynamic shapes

MLIR's type system makes the mean computation require a long chain of casts. I
checked that it all constant folds away for static shapes.

PiperOrigin-RevId: 373328525

integrations/tensorflow/iree_tf_compiler/TF/ConvertToMHLO.cpp

diff --git a/integrations/tensorflow/iree_tf_compiler/TF/ConvertToMHLO.cpp b/integrations/tensorflow/iree_tf_compiler/TF/ConvertToMHLO.cpp
index 7a3ab0f..8f03948 100644
--- a/integrations/tensorflow/iree_tf_compiler/TF/ConvertToMHLO.cpp
+++ b/integrations/tensorflow/iree_tf_compiler/TF/ConvertToMHLO.cpp
@@ -92,6 +92,7 @@
     target.addLegalDialect<mhlo::MhloDialect>();
     target.addLegalDialect<mlir::StandardOpsDialect>();
     target.addLegalDialect<shape::ShapeDialect>();
+    target.addLegalDialect<tensor::TensorDialect>();
     target.addLegalOp<mlir::CallOp>();
     target.addLegalOp<mlir::tensor::CastOp>();