Fix gather by using xla sparse gather (#3450)

All tests passing. Merging the PR.

Author: ymwangg

torch_xla/csrc/ops/gather.cpp

@@ -9,24 +9,11 @@
 namespace torch_xla {
 namespace ir {
 namespace ops {
-namespace {
-
-xla::Shape NodeOutputShape(const Value& input, const Value& index,
-                           int64_t dim) {
-  auto lower_for_shape_fn =
-      [&](absl::Span<const xla::XlaOp> operands) -> xla::XlaOp {
-    return xla::TorchGather(operands[0], operands[1], dim,
-                            IsSparseGather(operands[0], operands[1], dim));
-  };
-  return InferOutputShape({input.xla_shape(), index.xla_shape()},
-                          lower_for_shape_fn);
-}
-
-}  // namespace
 
 Gather::Gather(const Value& input, int64_t dim, const Value& index)
     : Node(torch::lazy::OpKind(at::aten::gather), {input, index},
-           [&]() { return NodeOutputShape(input, index, dim); },
+           xla::ShapeUtil::MakeShape(input.xla_shape().element_type(),
+                                     index.xla_shape().dimensions()),
            /*num_outputs=*/1, torch::lazy::MHash(dim)),
       dim_(dim) {}
 
@@ -37,9 +24,7 @@ NodePtr Gather::Clone(OpList operands) const {
 XlaOpVector Gather::Lower(LoweringContext* loctx) const {
   xla::XlaOp input = loctx->GetOutputOp(operand(0));
   xla::XlaOp index = loctx->GetOutputOp(operand(1));
-  return ReturnOp(
-      xla::TorchGather(input, index, dim_, IsSparseGather(input, index, dim_)),
-      loctx);
+  return ReturnOp(xla::TorchGather(input, index, dim_, /*sparse=*/true), loctx);
 }
 
 std::string Gather::ToString() const {

torch_xla/csrc/tensor_methods.cpp

@@ -1391,10 +1391,18 @@ XLATensor XLATensor::full_like(const XLATensor& input,
 
 XLATensor XLATensor::gather(const XLATensor& input, int64_t dim,
                             const XLATensor& index) {
+  xla::Shape input_shape = input.shape();
+  xla::Shape index_shape = index.shape();
+  XLA_CHECK_EQ(input_shape.rank(), index_shape.rank());
+  int64_t canonical_dim =
+      torch::lazy::GetCanonicalDimensionIndex(dim, input_shape.rank());
+  for (size_t dim = 0; dim < input_shape.rank(); dim++) {
+    if (dim != canonical_dim) {
+      XLA_CHECK_LE(index.size(dim), input.size(dim));
+    }
+  }
   return input.CreateFrom(ir::MakeNode<ir::ops::Gather>(
-      input.GetIrValue(),
-      torch::lazy::GetCanonicalDimensionIndex(dim, input.shape().get().rank()),
-      index.GetIrValue()));
+      input.GetIrValue(), canonical_dim, index.GetIrValue()));
 }
 
 XLATensor XLATensor::ge(const XLATensor& input, const at::Scalar& other) {