Fix torch.min and torch.mode (#7513) (#8092)

Co-authored-by: Simon Teo <simonteo@google.com>

experimental/torch_xla2/test/test_ops.py

@@ -69,8 +69,6 @@
     "lu_unpack",
     "masked.median",
     "max_pool2d_with_indices_backward",
-    "min",
-    "mode",
     "multinomial",
     "mvlgamma",
     "nanmedian",
@@ -243,7 +241,8 @@ def run_export_and_compare(testcase,
 # For example: sort( [1, 0, 0]) -> [0, 0, 1]
 # the correct index can be [1, 2, 0] or [2, 1, 0]
 should_ignore_indexes = {
-  "topk"
+  "topk",
+  "mode"
 }
 
 

experimental/torch_xla2/torch_xla2/ops/jaten.py

@@ -1079,11 +1079,25 @@ def reduce_fn(a, b):
 
 
 @op(torch.ops.aten.min)
-def _aten_min(x, axis=None):
-  if axis:
-    return jnp.min(x, axis=axis), jnp.argmin(x, axis=axis).astype(jnp.int64)
+def _aten_min(x, dim=None, keepdim=False):
+  if dim is not None:
+    return _with_reduction_scalar(jnp.min, x, dim, keepdim), _with_reduction_scalar(jnp.argmin, x, dim, keepdim).astype(jnp.int64)
   else:
-    return jnp.min(x, axis=axis)
+    return _with_reduction_scalar(jnp.min, x, dim, keepdim)
+
+
+@op(torch.ops.aten.mode)
+def _aten_mode(input, dim=-1, keepdim=False, *, out=None):
+  if input.ndim == 0: # single number
+    return input, jnp.array(0)
+  dim = (input.ndim + dim) % input.ndim # jnp.scipy.stats.mode does not accept -1 as dim
+  # 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)
+  if not keepdim:
+    mode = mode.squeeze(axis=dim)
+  indices = jnp.argmax(jnp.equal(mode_broadcast, input), axis=dim, keepdims=keepdim)
+  return mode, indices
 
 
 @op(torch.ops.aten.amin)