Add support for torch.unravel_index. (#8266)

Co-authored-by: mrguenther <[email protected]>
pytorch · Oct 18, 2024 · a3e4e12 · a3e4e12
1 parent d1a1966
commit a3e4e12
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -90,7 +90,6 @@
     "svd_lowrank",
     "unfold_copy",
     "unfold",
-    "unravel_index",
     "nanmean",
     "nn.functional.upsample_bilinear",
     "randint",

diff --git a/experimental/torch_xla2/torch_xla2/ops/jtorch.py b/experimental/torch_xla2/torch_xla2/ops/jtorch.py
@@ -245,6 +245,11 @@ def empty_strided(
   return empty(size, dtype=dtype)
 
 
+@register_function(torch.unravel_index)
+def unravel_index(indices, shape):
+  return jnp.unravel_index(indices, shape)
+
+
 @register_function(torch.rand, is_jax_function=False)
 def rand(
   *size, **kwargs