From 15e945a7910e9d639c9711fbdbdcf262eb9fe8e6 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Fri, 6 Mar 2020 02:34:32 -0800 Subject: [PATCH] Add additional elementwise unary ops to hlo_shape_derivation. This makes them work with dynamic shapes. PiperOrigin-RevId: 299311603 Change-Id: Id3801de8406ffc0054a7f6cf6c3ac25c4dd95f55 --- .../compiler/mlir/xla/transforms/hlo_shape_derivation.h | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tensorflow/compiler/mlir/xla/transforms/hlo_shape_derivation.h b/tensorflow/compiler/mlir/xla/transforms/hlo_shape_derivation.h index d2a1f47e540..0a989fea059 100644 --- a/tensorflow/compiler/mlir/xla/transforms/hlo_shape_derivation.h +++ b/tensorflow/compiler/mlir/xla/transforms/hlo_shape_derivation.h @@ -114,11 +114,16 @@ SAME_SHAPE_AS_FIRST_OPERAND(CeilOp) SAME_SHAPE_AS_FIRST_OPERAND(CosOp) SAME_SHAPE_AS_FIRST_OPERAND(DivOp) SAME_SHAPE_AS_FIRST_OPERAND(ExpOp) +SAME_SHAPE_AS_FIRST_OPERAND(LogOp) SAME_SHAPE_AS_FIRST_OPERAND(MaxOp) SAME_SHAPE_AS_FIRST_OPERAND(MinOp) SAME_SHAPE_AS_FIRST_OPERAND(MulOp) SAME_SHAPE_AS_FIRST_OPERAND(NegOp) SAME_SHAPE_AS_FIRST_OPERAND(RemOp) +SAME_SHAPE_AS_FIRST_OPERAND(RsqrtOp) +SAME_SHAPE_AS_FIRST_OPERAND(SignOp) +SAME_SHAPE_AS_FIRST_OPERAND(SinOp) +SAME_SHAPE_AS_FIRST_OPERAND(SqrtOp) SAME_SHAPE_AS_FIRST_OPERAND(SubOp) SAME_SHAPE_AS_FIRST_OPERAND(TanhOp)