Add leaky_relu_ to ATen XLA tensor

Author: asuhan

test/cpp/test_aten_xla_tensor.cpp

@@ -1715,6 +1715,18 @@ TEST_F(AtenXlaTensorTest, TestLeakyRelu) {
   });
 }
 
+TEST_F(AtenXlaTensorTest, TestLeakyReluInPlace) {
+  at::Tensor input = at::rand({2, 1, 4, 6}, at::TensorOptions(at::kFloat));
+  double negative_slope = 0.01;
+  ForEachDevice([&](const Device& device) {
+    at::Tensor xla_input = bridge::CreateXlaTensor(input.clone(), device);
+    at::Tensor output = at::leaky_relu_(input, negative_slope);
+    at::Tensor xla_output = at::leaky_relu_(xla_input, negative_slope);
+    AllClose(output, xla_output);
+    AllClose(input, xla_input);
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestExp) {
   at::Tensor a = at::rand({2, 2}, at::TensorOptions(at::kFloat));
   at::Tensor b = at::exp(a);

torch_xla/csrc/aten_xla_type.cpp

@@ -1147,6 +1147,13 @@ at::Tensor AtenXlaType::leaky_relu(const at::Tensor& self,
       bridge::GetXlaTensor(self), negative_slope.to<double>()));
 }
 
+at::Tensor& AtenXlaType::leaky_relu_(at::Tensor& self,
+                                     at::Scalar negative_slope) const {
+  XLATensor self_tensor = bridge::GetXlaTensor(self);
+  XLATensor::leaky_relu_(self_tensor, negative_slope.to<double>());
+  return self;
+}
+
 at::Tensor AtenXlaType::threshold(const at::Tensor& self, at::Scalar threshold,
                                   at::Scalar value) const {
   return bridge::AtenFromXlaTensor(XLATensor::threshold(

torch_xla/csrc/aten_xla_type.h

@@ -348,6 +348,8 @@ class AtenXlaType : public AtenXlaTypeBase {
 
   at::Tensor leaky_relu(const at::Tensor& self,
                         at::Scalar negative_slope) const override;
+  at::Tensor& leaky_relu_(at::Tensor& self,
+                          at::Scalar negative_slope) const override;
 
   at::Tensor threshold(const at::Tensor& self, at::Scalar threshold,
                        at::Scalar value) const override;

torch_xla/csrc/tensor.cpp

@@ -907,6 +907,11 @@ XLATensor XLATensor::leaky_relu(const XLATensor& input, double negative_slope) {
       ir::MakeNode<ir::ops::LeakyRelu>(input.GetIrValue(), negative_slope));
 }
 
+void XLATensor::leaky_relu_(XLATensor& input, double negative_slope) {
+  input.SetIrValue(
+      ir::MakeNode<ir::ops::LeakyRelu>(input.GetIrValue(), negative_slope));
+}
+
 XLATensor XLATensor::DispatchComparisonOp(c10::Symbol kind,
                                           const XLATensor& input,
                                           const at::Scalar& other) {

torch_xla/csrc/tensor.h

@@ -216,6 +216,7 @@ class XLATensor {
   static void relu_(XLATensor& input);
 
   static XLATensor leaky_relu(const XLATensor& input, double negative_slope);
+  static void leaky_relu_(XLATensor& input, double negative_slope);
 
   static XLATensor threshold(const XLATensor& input, float threshold,
                              float value);