Support logistic for vector data type.

PiperOrigin-RevId: 606812186
tensorflow · Feb 14, 2024 · 4aedbe2 · 4aedbe2
1 parent 537d106
commit 4aedbe2
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/mhlo/transforms/map_mhlo_to_scalar_op.h b/mhlo/transforms/map_mhlo_to_scalar_op.h
@@ -1081,7 +1081,12 @@ inline Value mapMhloOpToStdScalarOp<mhlo::LogisticOp>(
   Value expNegX = mapMhloOpToStdScalarOp<mhlo::ExpOp>(loc, resultTypes,
                                                       resultTypes, {{negX}}, b);
 
-  Value oneFloat = b->create<arith::ConstantOp>(loc, b->getF32FloatAttr(1.0));
+  Type type = getElementTypeOrSelf(resultTypes[0]);
+  Value oneFloat =
+      type.isa<ComplexType>()
+          ? b->create<arith::ConstantOp>(loc, b->getF32FloatAttr(1.0))
+          : getConstantOrSplat(b, loc, resultTypes[0],
+                               FloatAttr::get(type, 1.0f));
   Value one = mapConvertOpToStdScalarOp(loc, resultTypes, resultTypes,
                                         {oneFloat.getType()}, {{oneFloat}}, b);
   Value oneAddExprNegX = mapMhloOpToStdScalarOp<mhlo::AddOp>(