Lower Elu and EluBackward

Author: wonjoo-wj

torch_xla/csrc/elementwise.cpp

@@ -342,4 +342,36 @@ xla::XlaOp BuildLogSigmoidBackward(xla::XlaOp grad_output,
   return grad_output * (xla::Neg(max_deriv) - sign * (buffer - one) / buffer);
 }
 
+xla::XlaOp BuildElu(xla::XlaOp input, const at::Scalar& alpha,
+                         const at::Scalar& scale,
+                         const at::Scalar& input_scale) {
+  const xla::Shape& shape = XlaHelpers::ShapeOfXlaOp(input);
+  xla::XlaOp scaled_input = input * XlaHelpers::ScalarValue(input_scale, shape.element_type(), input.builder());
+  xla::XlaOp zero = xla::Zero(input.builder(), shape.element_type());
+  xla::XlaOp one = XlaHelpers::ScalarValue<float>(1.0, shape.element_type(),
+                                                  input.builder());
+  xla::XlaOp alpha_scalar = XlaHelpers::ScalarValue(alpha, shape.element_type(), input.builder());
+  xla::XlaOp scale_scalar = XlaHelpers::ScalarValue(scale, shape.element_type(), input.builder());
+  return xla::Select(xla::Le(input, zero), alpha_scalar * (xla::Exp(scaled_input) - one), input) * scale_scalar;
+
+
+  // XlaHelpers::ScalarValue(lower, shape.element_type(), input.builder());
+
+}
+
+xla::XlaOp BuildEluBackward(xla::XlaOp grad_output,
+                                 xla::XlaOp output,
+                                 const at::Scalar& alpha,
+                                 const at::Scalar& scale,
+                                 const at::Scalar& input_scale) {
+  const xla::Shape& shape = XlaHelpers::ShapeOfXlaOp(output);
+  xla::XlaOp zero = xla::Zero(output.builder(), shape.element_type());
+  xla::XlaOp alpha_scalar = XlaHelpers::ScalarValue(alpha, shape.element_type(), output.builder());
+  xla::XlaOp scale_scalar = XlaHelpers::ScalarValue(scale, shape.element_type(), output.builder());
+  xla::XlaOp input_scale_scalar = XlaHelpers::ScalarValue(input_scale, shape.element_type(), output.builder());
+  xla::XlaOp negative_output_branch = input_scale_scalar * (output + alpha_scalar * scale_scalar);
+  return grad_output * xla::Select(xla::Gt(output, zero), scale_scalar, negative_output_branch);
+  
+}
+
 }  // namespace torch_xla

torch_xla/csrc/elementwise.h

@@ -84,8 +84,21 @@ xla::XlaOp BuildGeluBackward(xla::XlaOp grad_output, xla::XlaOp input);
 // Computes the LogSigmoid function of input.
 std::vector<xla::XlaOp> BuildLogSigmoid(xla::XlaOp input);
 
+// Computes the backward of LogSigmoid.
 xla::XlaOp BuildLogSigmoidBackward(xla::XlaOp grad_output,
                                         xla::XlaOp input,
                                         xla::XlaOp buffer);
 
+// Computes the Elu function of input.
+xla::XlaOp BuildElu(xla::XlaOp input, const at::Scalar& alpha,
+                         const at::Scalar& scale,
+                         const at::Scalar& input_scale);
+
+// Computes the backward of Elu.
+xla::XlaOp BuildEluBackward(xla::XlaOp grad_output,
+                                 xla::XlaOp output,
+                                 const at::Scalar& alpha,
+                                 const at::Scalar& scale,
+                                 const at::Scalar& input_scale);
+
 }  // namespace torch_xla

torch_xla/csrc/ops/ops.cpp

@@ -709,31 +709,30 @@ torch::lazy::NodePtr Identity(int64_t lines, int64_t cols,
 torch::lazy::NodePtr Elu(const XlaValue& input, const at::Scalar& alpha,
                          const at::Scalar& scale,
                          const at::Scalar& input_scale) {
-  ScopePusher ir_scope(at::aten::elu.toQualString());
-  const xla::Shape& shape = input.xla_shape();
-  torch::lazy::NodePtr scaled_input = input * ScalarOp(input_scale, shape);
-  torch::lazy::NodePtr zero = ScalarOp(0, shape);
-  torch::lazy::NodePtr one = ScalarOp(1, shape);
-  torch::lazy::NodePtr alpha_scalar = ScalarOp(alpha, shape);
-  return Where(ComparisonOp(at::aten::le, input, zero),
-               alpha_scalar * (Exp(scaled_input) - one), input) *
-         ScalarOp(scale, shape);
+  auto lower_fn = [=](const XlaNode& node,
+                     LoweringContext* loctx) -> XlaOpVector {
+    xla::XlaOp xla_input = loctx->GetOutputOp(node.operand(0));
+    return node.ReturnOp(BuildElu(xla_input, alpha, scale, input_scale), loctx);
+  };
+  return GenericOp(torch::lazy::OpKind(at::aten::elu),
+                   {input}, input.xla_shape(),
+                   std::move(lower_fn));
 }
 
 torch::lazy::NodePtr EluBackward(const XlaValue& grad_output,
                                  const XlaValue& output,
                                  const at::Scalar& alpha,
                                  const at::Scalar& scale,
                                  const at::Scalar& input_scale) {
-  ScopePusher ir_scope(at::aten::elu_backward.toQualString());
-  const xla::Shape& shape = grad_output.xla_shape();
-  torch::lazy::NodePtr negative_output_branch =
-      ScalarOp(input_scale, shape) *
-      (output + ScalarOp(alpha, shape) * ScalarOp(scale, shape));
-  torch::lazy::NodePtr positive_output_branch = ScalarOp(scale, shape);
-  return grad_output *
-         Where(ComparisonOp(at::aten::gt, output, ScalarOp(0, shape)),
-               positive_output_branch, negative_output_branch);
+  auto lower_fn = [=](const XlaNode& node,
+                     LoweringContext* loctx) -> XlaOpVector {
+    xla::XlaOp xla_grad_output = loctx->GetOutputOp(node.operand(0));
+    xla::XlaOp xla_output = loctx->GetOutputOp(node.operand(1));
+    return node.ReturnOp(BuildEluBackward(xla_grad_output, xla_output, alpha, scale, input_scale), loctx);
+  };
+  return GenericOp(torch::lazy::OpKind(at::aten::elu_backward),
+                   {grad_output, output}, output.xla_shape(),
+                   std::move(lower_fn));
 }
 
 torch::lazy::NodePtr Gelu(const XlaValue& input) {