Fix torch digamma (#7383) (#8260)

Co-authored-by: Simon Teo <[email protected]>

experimental/torch_xla2/test/test_ops.py

@@ -18,7 +18,6 @@
     "cholesky",
     "cholesky_solve",
     "diagonal_copy",
-    "digamma",
     "geqrf",
     "histogram", # hard op: AssertionError: Tensor-likes are not close!
     "histogramdd", # TypeError: histogram requires ndarray or scalar arguments, got <class 'list'> at position 1.

experimental/torch_xla2/torch_xla2/ops/jaten.py

@@ -2394,6 +2394,12 @@ def _aten_hypot(input, other):
   return jnp.hypot(input, other)
 
 
+@op(torch.ops.aten.digamma)
+def _aten_digamma(input, *, out=None):
+  res = jax.scipy.special.digamma(input).astype(jnp.float32)
+  # replace indices where input == 0 with -inf in res
+  return jnp.where(jnp.equal(input, jnp.zeros(input.shape)), -jnp.inf, res)
+
 @op(torch.ops.aten.igamma)
 def _aten_igamma(input, other):
   return jax.scipy.special.gammainc(input, other)