Fix torch.mode tests

experimental/torch_xla2/test/test_ops.py

@@ -73,8 +73,6 @@
     "lu_unpack",
     "masked.median",
     "max_pool2d_with_indices_backward",
-    "min",
-    "mode",
     "multinomial",
     "mvlgamma",
     "nanmedian",

experimental/torch_xla2/torch_xla2/ops/jaten.py

@@ -1067,6 +1067,28 @@ def _aten_min(x, dim=None, keepdim=False):
     return _with_reduction_scalar(jnp.min, x, dim, keepdim)
 
 
+@op(torch.ops.aten.mode)
+def _aten_mode(input, dim=-1, keepdim=False, *, out=None):
+  print("INPUT SHAPE: ", input.shape)
+  if input.ndim == 0:
+    keepdim = False # no dimensions to keep
+    dim = 0
+    input = jnp.expand_dims(input, 0)
+  else:
+    dim = (input.ndim + dim) % input.ndim
+  # keepdims must be True for accurate broadcasting
+  mode, _ = jax.scipy.stats.mode(input, axis=dim, keepdims=True)
+  mode_broadcast = jnp.broadcast_to(mode, input.shape)
+  # find last occurence of mode value
+  flip_input =  jnp.flip(input, axis=dim)
+  indices = jnp.argmax(jnp.equal(mode_broadcast, flip_input), axis=dim, keepdims=keepdim)
+  len_array = jnp.ones(indices.shape).astype(jnp.int64) * (input.shape[dim] - 1)
+  res_indices = len_array - indices
+  if not keepdim:
+    mode, _ = jax.scipy.stats.mode(input, axis=dim, keepdims=False)
+  return mode, res_indices
+
+
 @op(torch.ops.aten.amin)
 def _aten_amin(x, dim=None, keepdim=False):
   return _with_reduction_scalar(jnp.amin, x, dim, keepdim)