Added aten::sigmoid operations.

Author: dlibenzi

test/cpp/test_aten_xla_tensor.cpp

@@ -584,6 +584,16 @@ TEST_F(AtenXlaTensorTest, TestAbs) {
   });
 }
 
+TEST_F(AtenXlaTensorTest, TestSigmoid) {
+  at::Tensor a = at::rand({2, 2}, at::TensorOptions(at::kFloat));
+  at::Tensor b = at::sigmoid(a);
+  ForEachDevice([&](const Device& device) {
+    at::Tensor xla_a = bridge::CreateXlaTensor(a, device);
+    at::Tensor xla_b = at::sigmoid(xla_a);
+    AllClose(b, xla_b, /*rtol=*/1e-3, /*atol=*/1e-5);
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestAddCMul) {
   at::Tensor a = at::rand({2, 2}, at::TensorOptions(at::kFloat));
   at::Tensor b = at::rand({2, 2}, at::TensorOptions(at::kFloat));

torch_xla/csrc/aten_xla_type.cpp

@@ -541,6 +541,17 @@ at::Tensor AtenXlaType::softmax(const at::Tensor& self, int64_t dim) const {
       XLATensor::softmax(bridge::GetXlaTensor(self), dim));
 }
 
+at::Tensor AtenXlaType::sigmoid(const at::Tensor& self) const {
+  return bridge::AtenFromXlaTensor(
+      XLATensor::sigmoid(bridge::GetXlaTensor(self)));
+}
+
+at::Tensor& AtenXlaType::sigmoid_(at::Tensor& self) const {
+  XLATensor self_tensor = bridge::GetXlaTensor(self);
+  XLATensor::sigmoid_(self_tensor);
+  return self;
+}
+
 at::Tensor AtenXlaType::max_pool2d(const at::Tensor& self,
                                    at::IntList kernel_size, at::IntList stride,
                                    at::IntList padding, at::IntList dilation,

torch_xla/csrc/aten_xla_type.h

@@ -184,6 +184,9 @@ class AtenXlaType : public AtenXlaTypeBase {
 
   at::Tensor softmax(const at::Tensor& self, int64_t dim) const override;
 
+  at::Tensor sigmoid(const at::Tensor& self) const override;
+  at::Tensor& sigmoid_(at::Tensor& self) const override;
+
   at::Tensor max_pool2d(const at::Tensor& self, at::IntList kernel_size,
                         at::IntList stride, at::IntList padding,
                         at::IntList dilation, bool ceil_mode) const override;

torch_xla/csrc/elementwise.cpp

@@ -91,4 +91,11 @@ xla::XlaOp BuildTypeAs(const torch::jit::Node* node,
   return xla::ConvertElementType(operand, target_type);
 }
 
+xla::XlaOp BuildSigmoid(const xla::XlaOp& input) {
+  xla::Shape shape = XlaHelpers::ShapeOfXlaOp(input);
+  xla::XlaOp half =
+      XlaHelpers::ScalarValue<float>(0.5, shape.element_type(), input.builder());
+  return  half + half * xla::Tanh(half * input);
+}
+
 }  // namespace torch_xla

torch_xla/csrc/elementwise.h

@@ -28,4 +28,6 @@ xla::XlaOp BuildThreshold(const xla::XlaOp& input, const xla::XlaOp& output,
 // Computes the rectified linear unit (replace negative elements with 0).
 xla::XlaOp BuildRelu(const xla::XlaOp& input);
 
+xla::XlaOp BuildSigmoid(const xla::XlaOp& input);
+
 }  // namespace torch_xla

torch_xla/csrc/ops/ops.cpp

@@ -90,6 +90,16 @@ NodePtr TransposeOp(const Value& input) {
                             output_shape, std::move(lower_fn));
 }
 
+NodePtr Sigmoid(const Value& input) {
+  auto lower_fn = [](const ir::Node& node,
+                     ir::LoweringContext* loctx) -> ir::XlaOpVector {
+    xla::XlaOp xla_input = loctx->GetOutputOp(node.operand(0));
+    return node.ReturnOp(BuildSigmoid(xla_input), loctx);
+  };
+  return ir::ops::GenericOp(ir::OpKind(at::aten::sigmoid), ir::OpList{input},
+                            input.shape(), std::move(lower_fn));
+}
+
 NodePtr Clamp(const Value& input, c10::optional<at::Scalar> min,
               c10::optional<at::Scalar> max) {
   const xla::Shape& input_shape = input.shape();

torch_xla/csrc/ops/ops.h

@@ -81,6 +81,8 @@ NodePtr Sqrt(const Value& input);
 
 NodePtr Pow(const Value& input, const Value& exponent);
 
+NodePtr Sigmoid(const Value& input);
+
 NodePtr Clamp(const Value& input, c10::optional<at::Scalar> min,
               c10::optional<at::Scalar> max);
 

torch_xla/csrc/tensor.cpp

@@ -1136,6 +1136,14 @@ XLATensor XLATensor::softmax(const XLATensor& input, xla::int64 dim) {
                 input.GetDevice());
 }
 
+XLATensor XLATensor::sigmoid(const XLATensor& input) {
+  return Create(ir::ops::Sigmoid(input.GetIrValue()), input.GetDevice());
+}
+
+void XLATensor::sigmoid_(XLATensor& input) {
+  input.SetIrValue(ir::ops::Sigmoid(input.GetIrValue()));
+}
+
 XLATensor XLATensor::nll_loss(const XLATensor& input, const XLATensor& target) {
   return Create(ir::ops::NllLossOp(input.GetIrValue(), target.GetIrValue()),
                 input.GetDevice());

torch_xla/csrc/tensor.h

@@ -196,6 +196,9 @@ class XLATensor {
 
   static XLATensor softmax(const XLATensor& input, xla::int64 dim);
 
+  static XLATensor sigmoid(const XLATensor& input);
+  static void sigmoid_(XLATensor& input);
+
   static XLATensor ones(tensorflow::gtl::ArraySlice<const xla::int64> size,
                         const Device& device, at::ScalarType scalar_type);
   static XLATensor ones_like(const XLATensor& input, const Device& device,

torch_xla/csrc/translator.cpp

@@ -366,11 +366,7 @@ void TranslateSigmoid(const torch::jit::Node* node, ComputationContext* cctx,
                       xla::XlaBuilder* b) {
   XLA_CHECK_EQ(node->inputs().size(), 1);
   xla::XlaOp xla_input = cctx->OpForInput(node, 0);
-  xla::Shape xla_input_shape = XlaHelpers::ShapeOfXlaOp(xla_input);
-  xla::XlaOp half =
-      XlaHelpers::ScalarValue<float>(0.5, xla_input_shape.element_type(), b);
-  xla::XlaOp xla_output = half + half * xla::Tanh(half * xla_input);
-  cctx->AddNodeOp(node, xla_output);
+  cctx->AddNodeOp(node, BuildSigmoid(xla_input));
 }
 
 void TranslateRelu(const torch::jit::Node* node, ComputationContext* cctx,