Fix quantile, nanquantile and gather (#8040)

pytorch · Sep 18, 2024 · 6b33f8f · 6b33f8f
1 parent 0415cc6
commit 6b33f8f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -32,7 +32,6 @@
     "digamma",
     "erfinv",
     "exponential",
-    "gather",
     "gcd",
     "geometric",
     "geqrf",
@@ -88,7 +87,6 @@
     "multinomial",
     "mvlgamma",
     "nanmedian",
-    "nanquantile",
     "new_empty",
     "new_empty_strided",
     "nextafter",
@@ -148,7 +146,6 @@
     "polygamma",
     "prod",
     "put",
-    "quantile",
     "repeat_interleave",
     "resize_",
     "resize_as_",

diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -40,6 +40,10 @@
   torch.ops.aten.bernoulli_: torch.ops.aten.bernoulli.p,
   torch.ops.aten.clamp_: torch.ops.aten.clamp,
   torch.ops.aten.random_: torch.ops.aten.uniform,
+  torch.ops.aten.ceil_: torch.ops.aten.ceil,
+  torch.ops.aten.logical_not_: torch.ops.aten.logical_not,
+  torch.ops.aten.unsqueeze_: torch.ops.aten.unsqueeze,
+  torch.ops.aten.transpose_: torch.ops.aten.transpose,
 }
 
 
@@ -1476,11 +1480,6 @@ def _aten_pixel_shuffle(x, upscale_factor):
 def _aten_lt(self, other):
   return self < other
 
-# aten.logical_not_
-@op(torch.ops.aten.logical_not_)
-def _aten_logical_not_(input):
-  return jnp.logical_not(input)
-
 
 def pool(inputs, init, reduce_fn, window_shape, strides, padding):
   """Helper function to define pooling functions.
@@ -2043,6 +2042,10 @@ def _aten_frexp(input):
 # aten.gather
 @op(torch.ops.aten.gather)
 def _aten_gather(input, dim, index):
+  if input.ndim == 0:
+    return jnp.broadcast_to(input, index.shape)
+  if dim < 0:
+    dim += input.ndim
   input_indexes, source_indexes = _scatter_index(dim, index)
   return input[input_indexes]