Codegen hardshrink (#3999)

* compiles

* final cleanup

* reformat cpp files.

* fixed a typo

* Do the necessary type conversion.

Author: vanbasten23

test/cpp/test_aten_xla_tensor.cpp

@@ -5854,6 +5854,17 @@ TEST_F(AtenXlaTensorTest, TestHardshrink) {
   });
 }
 
+TEST_F(AtenXlaTensorTest, TestHardshrinkWithNegativeLambda) {
+  torch::Tensor input = torch::randn({10}, torch::TensorOptions(torch::kFloat));
+  torch::Scalar lambd = -0.5;
+  torch::Tensor output = torch::hardshrink(input, lambd);
+  ForEachDevice([&](const torch::Device& device) {
+    torch::Tensor xla_input = CopyToDevice(input, device);
+    torch::Tensor xla_output = torch::hardshrink(xla_input, lambd);
+    AllClose(output, xla_output);
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestHardSigmoid) {
   torch::Tensor input = torch::randn({10}, torch::TensorOptions(torch::kFloat));
   torch::Tensor output = torch::hardsigmoid(input);

torch_xla/csrc/aten_xla_type.cpp

@@ -1246,13 +1246,6 @@ at::Tensor XLANativeFunctions::gelu_backward(const at::Tensor& grad,
       bridge::GetXlaTensor(grad), bridge::GetXlaTensor(self), approximate));
 }
 
-at::Tensor XLANativeFunctions::hardshrink(const at::Tensor& self,
-                                          const at::Scalar& lambda) {
-  XLA_FN_COUNTER("xla::");
-  return bridge::AtenFromXlaTensor(
-      XLATensor::hardshrink(bridge::GetXlaTensor(self), lambda));
-}
-
 at::Tensor XLANativeFunctions::hardshrink_backward(const at::Tensor& grad_out,
                                                    const at::Tensor& self,
                                                    const at::Scalar& lambda) {

torch_xla/csrc/elementwise.cpp

@@ -69,10 +69,20 @@ xla::XlaOp BuildRelu(xla::XlaOp input) {
                              0, input_shape.element_type(), input.builder()));
 }
 
-xla::XlaOp BuildHardshrink(xla::XlaOp input, const at::Scalar& lambda) {
+xla::XlaOp BuildHardshrink(xla::XlaOp input, xla::XlaOp lambda) {
   const xla::Shape& shape = XlaHelpers::ShapeOfXlaOp(input);
-  xla::XlaOp zero = xla::Zero(input.builder(), shape.element_type());
-  return xla::Select(Between(input, -lambda, lambda), zero, input);
+  xla::PrimitiveType input_element_type = shape.element_type();
+  xla::XlaOp zero = xla::Zero(input.builder(), input_element_type);
+
+  // The conversion here is needed because when we do computation such as
+  // broadcast or subtraction for input and lambda, XLA disallows mixed
+  // precision for float point types.
+  lambda = MaybeConvertTo(lambda, input_element_type);
+  xla::XlaOp check_low = BuildComparisonOp(at::aten::ge, input, zero - lambda);
+  xla::XlaOp check_high = BuildComparisonOp(at::aten::le, input, lambda);
+  xla::XlaOp between = xla::And(check_low, check_high);
+
+  return xla::Select(between, zero, input);
 }
 
 xla::XlaOp BuildHardSigmoid(xla::XlaOp input) {

torch_xla/csrc/elementwise.h

@@ -28,7 +28,7 @@ xla::XlaOp BuildRreluBackward(xla::XlaOp grad_output, xla::XlaOp input,
                               xla::XlaOp noise, const at::Scalar& lower,
                               const at::Scalar& upper, bool training);
 
-xla::XlaOp BuildHardshrink(xla::XlaOp input, const at::Scalar& lambda);
+xla::XlaOp BuildHardshrink(xla::XlaOp input, xla::XlaOp lambda);
 xla::XlaOp BuildHardSigmoid(xla::XlaOp input);
 xla::XlaOp BuildHardSigmoidBackward(xla::XlaOp grad_output, xla::XlaOp input);
 xla::XlaOp BuildHardSwish(xla::XlaOp input);

torch_xla/csrc/ops/hardshrink.cpp

@@ -330,6 +330,12 @@ torch_xla::XlaOpVector GtTensor::Lower(LoweringContext* loctx) const {
   return ReturnOp(BuildComparisonOp(at::aten::gt, xla_input, xla_other), loctx);
 }
 
+torch_xla::XlaOpVector Hardshrink::Lower(LoweringContext* loctx) const {
+  xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
+  xla::XlaOp lambd = loctx->GetOutputOp(operand(1));
+  return ReturnOp(BuildHardshrink(xla_input, lambd), loctx);
+}
+
 torch_xla::XlaOpVector Hardsigmoid::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
   return ReturnOp(BuildHardSigmoid(xla_input), loctx);

torch_xla/csrc/ops/hardshrink.h

@@ -409,6 +409,11 @@ xla::Shape GtTensorOutputShape(const torch::lazy::Value& self,
   return GtScalarOutputShape(self, other);
 }
 
+xla::Shape HardshrinkOutputShape(const torch::lazy::Value& self,
+                                 const torch::lazy::Value& lambd) {
+  return GetXlaShape(self);
+}
+
 xla::Shape HardsigmoidOutputShape(const torch::lazy::Value& input) {
   return GetXlaShape(input);
 }

torch_xla/csrc/ops/ops_lower_fn.cpp

@@ -131,6 +131,9 @@ xla::Shape GtScalarOutputShape(const torch::lazy::Value& self,
 xla::Shape GtTensorOutputShape(const torch::lazy::Value& self,
                                const torch::lazy::Value& other);
 
+xla::Shape HardshrinkOutputShape(const torch::lazy::Value& self,
+                                 const torch::lazy::Value& lambd);
+
 xla::Shape HardsigmoidOutputShape(const torch::lazy::Value& input);
 
 xla::Shape HardsigmoidBackwardOutputShape(const torch::lazy::Value& grad_output,

torch_xla/csrc/ops/ops_xla_shape_fn.cpp

@@ -52,7 +52,6 @@
 #include "torch_xla/csrc/ops/gather.h"
 #include "torch_xla/csrc/ops/generic.h"
 #include "torch_xla/csrc/ops/get_dimensions_size.h"
-#include "torch_xla/csrc/ops/hardshrink.h"
 #include "torch_xla/csrc/ops/hardtanh_backward.h"
 #include "torch_xla/csrc/ops/index_ops.h"
 #include "torch_xla/csrc/ops/index_select.h"
@@ -1349,12 +1348,6 @@ XLATensorPtr XLATensor::le(const XLATensorPtr& input,
   return DispatchComparisonOp(at::aten::le, input, other);
 }
 
-XLATensorPtr XLATensor::hardshrink(const XLATensorPtr& input,
-                                   const at::Scalar& lambda) {
-  return input->CreateFrom(
-      torch::lazy::MakeNode<Hardshrink>(input->GetIrValue(), lambda));
-}
-
 XLATensorPtr XLATensor::hardshrink_backward(const XLATensorPtr& grad_out,
                                             const XLATensorPtr& input,
                                             const at::Scalar& lambda) {