Add selu_ back

Author: JackCaoG

test/cpp/test_aten_xla_tensor.cpp

@@ -6382,14 +6382,12 @@ TEST_F(AtenXlaTensorTest, TestSeluInPlace) {
     torch::Tensor xla_input = CopyToDevice(input, device);
     torch::Tensor output = torch::selu_(input);
     torch::Tensor xla_output = torch::selu_(xla_input);
-    std::cerr << output << "\n";
-    std::cerr << xla_output << "\n";
     AllClose(output, xla_output);
     AllClose(input, xla_input);
   });
 
-  // selu_ uses elu_ instead of selu
   ExpectCounterNotChanged("aten::.*", cpp_test::GetIgnoredCounters());
+  ExpectCounterChanged("xla::selu_", cpp_test::GetIgnoredCounters());
 }
 
 TEST_F(AtenXlaTensorTest, TestCelu) {

torch_xla/csrc/aten_xla_type.cpp

@@ -2785,6 +2785,14 @@ at::Tensor XLANativeFunctions::select(const at::Tensor& self, int64_t dim,
       XLATensor::select(bridge::GetXlaTensor(self), dim, index));
 }
 
+// TODO(JackCaoG): Remove after elu being codegened
+at::Tensor& XLANativeFunctions::selu_(at::Tensor& self) {
+  XLA_FN_COUNTER("xla::");
+  XLATensorPtr self_tensor = bridge::GetXlaTensor(self);
+  XLATensor::selu_(self_tensor);
+  return self;
+}
+
 at::Tensor XLANativeFunctions::sigmoid(const at::Tensor& self) {
   XLA_FN_COUNTER("xla::");
   return bridge::AtenFromXlaTensor(

torch_xla/csrc/tensor.h

@@ -1037,7 +1037,6 @@ class XLATensor : public c10::intrusive_ptr_target {
   static XLATensorPtr select(const XLATensorPtr& input, int64_t dim,
                              int64_t index);
 
-  static XLATensorPtr selu(const XLATensorPtr& input);
   static void selu_(XLATensorPtr& input);
 
   static XLATensorPtr silu(const XLATensorPtr& input);

torch_xla/csrc/tensor_methods.cpp

@@ -2420,6 +2420,10 @@ XLATensorPtr XLATensor::select(const XLATensorPtr& input, int64_t dim,
   return tensor_ops::Select(input, dim, index);
 }
 
+void XLATensor::selu_(XLATensorPtr& input) {
+  input->SetInPlaceIrValue(Selu(input->GetIrValue()));
+}
+
 XLATensorPtr XLATensor::sigmoid(const XLATensorPtr& input) {
   return input->CreateFrom(Sigmoid(input->GetIrValue()));
 }

xla_native_functions.yaml

@@ -274,6 +274,7 @@ supported:
   - scatter.value_reduce
   - scatter_add
   - select.int
+  - selu_
   - sigmoid
   - sigmoid_backward
   - slice.Tensor