implement diagonal_copy (#9416)

Signed-off-by: Chenyaaang <[email protected]>

Author: Chenyaaang

torchax/test/test_ops.py

@@ -14,7 +14,6 @@
     "bincount",  # NOTE: dtype for int input torch gives float. This is weird.
     "byte",
     "cholesky_solve",
-    "diagonal_copy",
     "geqrf",
     "histogram",  # hard op: AssertionError: Tensor-likes are not close!
     "histogramdd",  # TypeError: histogram requires ndarray or scalar arguments, got <class 'list'> at position 1.

torchax/torchax/ops/jaten.py

@@ -2495,6 +2495,7 @@ def _aten_diag(input, diagonal=0):
 
 # aten.diagonal
 @op(torch.ops.aten.diagonal)
+@op(torch.ops.aten.diagonal_copy)
 def _aten_diagonal(input, offset=0, dim1=0, dim2=1):
   return jnp.diagonal(input, offset, dim1, dim2)