Codegen take

torch_xla/csrc/aten_xla_type.cpp

@@ -3032,13 +3032,6 @@ at::Tensor& XLANativeFunctions::t_(at::Tensor& self) {
   return self;
 }
 
-at::Tensor XLANativeFunctions::take(const at::Tensor& self,
-                                    const at::Tensor& index) {
-  XLA_FN_COUNTER("xla::");
-  return bridge::AtenFromXlaTensor(
-      XLATensor::take(bridge::GetXlaTensor(self), bridge::GetXlaTensor(index)));
-}
-
 at::Tensor XLANativeFunctions::tanh_backward(const at::Tensor& grad_output,
                                              const at::Tensor& output) {
   XLA_FN_COUNTER("xla::");

torch_xla/csrc/ops/ops.cpp

@@ -713,21 +713,6 @@ torch::lazy::NodePtr MinUnary(const torch::lazy::Value& input) {
       std::move(lower_fn));
 }
 
-torch::lazy::NodePtr Take(const torch::lazy::Value& input,
-                          const torch::lazy::Value& index) {
-  auto lower_fn = [](const XlaNode& node,
-                     LoweringContext* loctx) -> XlaOpVector {
-    xla::XlaOp xla_input = loctx->GetOutputOp(node.operand(0));
-    xla::XlaOp xla_index = loctx->GetOutputOp(node.operand(1));
-    xla::XlaOp result = BuildTake(xla_input, xla_index);
-    return node.ReturnOp(result, loctx);
-  };
-  xla::Shape result_shape = GetXlaShape(index);
-  result_shape.set_element_type(GetXlaShape(input).element_type());
-  return GenericOp(torch::lazy::OpKind(at::aten::take), {input, index},
-                   std::move(result_shape), std::move(lower_fn));
-}
-
 torch::lazy::NodePtr TanhGelu(const torch::lazy::Value& input) {
   // TODO: add proper lowering function
   torch::lazy::ScopePusher ir_scope("aten::tanh_gelu");

torch_xla/csrc/ops/ops.h

@@ -215,9 +215,6 @@ torch::lazy::NodePtr MaxUnary(const torch::lazy::Value& input);
 
 torch::lazy::NodePtr MinUnary(const torch::lazy::Value& input);
 
-torch::lazy::NodePtr Take(const torch::lazy::Value& input,
-                          const torch::lazy::Value& index);
-
 torch::lazy::NodePtr TanhGelu(const torch::lazy::Value& input);
 
 torch::lazy::NodePtr TanhGeluBackward(const torch::lazy::Value& grad,

torch_xla/csrc/ops/ops_lower_fn.cpp

@@ -2,6 +2,7 @@
 
 #include "tensorflow/compiler/xla/client/lib/logdet.h"
 #include "tensorflow/compiler/xla/client/lib/math.h"
+#include "torch_xla/csrc/data_ops.h"
 #include "torch_xla/csrc/elementwise.h"
 #include "torch_xla/csrc/helpers.h"
 #include "torch_xla/csrc/matrix.h"
@@ -298,6 +299,13 @@ torch_xla::XlaOpVector Sinh::Lower(LoweringContext* loctx) const {
 //   return ReturnOps({result.sign, result.logdet}, loctx);
 // }
 
+torch_xla::XlaOpVector Take::Lower(LoweringContext* loctx) const {
+  xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
+  xla::XlaOp xla_index = loctx->GetOutputOp(operand(1));
+  xla::XlaOp result = BuildTake(xla_input, xla_index);
+  return ReturnOp(result, loctx);
+}
+
 torch_xla::XlaOpVector Tan::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
   return ReturnOp(xla::Tan(xla_input), loctx);

torch_xla/csrc/ops/ops_xla_shape_fn.cpp

@@ -330,6 +330,13 @@ xla::Shape TanOutputShape(const torch::lazy::Value& input) {
   return GetXlaShape(input);
 }
 
+xla::Shape TakeOutputShape(const torch::lazy::Value& input,
+                           const torch::lazy::Value& index) {
+  xla::Shape result_shape = GetXlaShape(index);
+  result_shape.set_element_type(GetXlaShape(input).element_type());
+  return result_shape;
+}
+
 xla::Shape TanhOutputShape(const torch::lazy::Value& input) {
   return GetXlaShape(input);
 }

torch_xla/csrc/ops/ops_xla_shape_fn.h

@@ -115,6 +115,9 @@ xla::Shape SinhOutputShape(const torch::lazy::Value& input);
 /* Blocked on https://github.com/pytorch/xla/issues/3596 */
 // xla::Shape SlogdetOutputShape(const torch::lazy::Value& input);
 
+xla::Shape TakeOutputShape(const torch::lazy::Value& input,
+                           const torch::lazy::Value& index);
+
 xla::Shape TanOutputShape(const torch::lazy::Value& input);
 
 xla::Shape TanhOutputShape(const torch::lazy::Value& input);

torch_xla/csrc/tensor_methods.cpp

@@ -2648,11 +2648,6 @@ std::tuple<XLATensorPtr, XLATensorPtr> XLATensor::symeig(
                          input->CreateFrom(torch::lazy::Value(node, 1)));
 }
 
-XLATensorPtr XLATensor::take(const XLATensorPtr& input,
-                             const XLATensorPtr& index) {
-  return input->CreateFrom(Take(input->GetIrValue(), index->GetIrValue()));
-}
-
 XLATensorPtr XLATensor::tanh_backward(const XLATensorPtr& grad_output,
                                       const XLATensorPtr& output) {
   return XLATensor::mul(grad_output,

xla_native_functions.yaml

@@ -45,6 +45,7 @@ full_codegen:
   - silu_backward
   - sin
   - sinh
+  - take
   - tan
   - tanh
   - tril
@@ -307,7 +308,6 @@ supported:
   - symeig
   - t
   - t_
-  - take
   - tanh_backward
   - threshold
   - threshold_backward