add histogram .. linalg.eigvals

experimental/torch_xla2/test/test_ops.py

@@ -28,7 +28,6 @@
     "complex",
     "cummax",
     "cummin",
-    "cumsum",
     "diag_embed",
     "diagflat",
     "diagonal_copy",
@@ -64,22 +63,18 @@
     "geometric",
     "geqrf",
     "grid_sampler_2d",
-    "histogram",
-    "histogramdd",
-    "hypot",
-    "igamma",
+    "histogram", # hard op: AssertionError: Tensor-likes are not close!
+    "histogramdd", # TypeError: histogram requires ndarray or scalar arguments, got <class 'list'> at position 1.
     "igammac",
     "index_reduce",
     "kthvalue",
     "lgamma",
     "linalg.cholesky",
     "linalg.cholesky_ex",
     "linalg.cond",
-    "linalg.cross",
     "linalg.det",
     "linalg.eig",
     "linalg.eigh",
-    "linalg.eigvals",
     "linalg.eigvalsh",
     "linalg.householder_product",
     "linalg.inv",

experimental/torch_xla2/torch_xla2/ops/jaten.py

@@ -1979,6 +1979,7 @@ def _aten_hardtanh(input, min_val=-1, max_val=1, inplace=False):
 # aten.histc
 @op(torch.ops.aten.histc)
 def _aten_histc(input, bins=100, min=0, max=0):
+  # TODO(@manfei): this function might cause some uncertainty
   if min==0 and max==0:
     if isinstance(input, jnp.ndarray) and input.size == 0:
       min = 0
@@ -1987,11 +1988,25 @@ def _aten_histc(input, bins=100, min=0, max=0):
       min = jnp.min(input)
       max = jnp.max(input)
   range_value = (min, max)
-  print("range_value: ", range_value)
   hist, bin_edges = jnp.histogram(input, bins=bins, range=range_value, weights=None, density=None)
   return hist
 
 
+@op(torch.ops.aten.hypot)
+def _aten_hypot(input, other):
+  return jnp.hypot(input, other)
+
+
+@op(torch.ops.aten.igamma)
+def _aten_igamma(input, other):
+  return jax.scipy.special.gammainc(input, other)
+
+
+@op(torch.ops.aten.linalg_eig)
+def _aten_linalg_eig(A):
+  return jax.numpy.linalg.eig(A)
+
+
 # aten.lcm
 @op(torch.ops.aten.lcm)
 def _aten_lcm(input, other):