From 4325c55b46e48a13a8289e12838b50fc9578a1a3 Mon Sep 17 00:00:00 2001
From: Irina Demeshko <idemeshko@nvidia.com>
Date: Wed, 22 Nov 2023 22:23:02 -0800
Subject: [PATCH] removing unused variable in diagonal implementation

---
 cunumeric/array.py  | 19 +++----------------
 cunumeric/module.py |  9 +++------
 2 files changed, 6 insertions(+), 22 deletions(-)

diff --git a/cunumeric/array.py b/cunumeric/array.py
index a5145269d..0612417ec 100644
--- a/cunumeric/array.py
+++ b/cunumeric/array.py
@@ -2535,10 +2535,9 @@ def _diag_helper(
     def diagonal(
         self,
         offset: int = 0,
-        axis1: Any = None,
-        axis2: Any = None,
+        axis1: int = 0,
+        axis2: int = 1,
         extract: bool = True,
-        axes: Any = None,
     ) -> ndarray:
         """a.diagonal(offset=0, axis1=None, axis2=None)
 
@@ -2560,19 +2559,7 @@ def diagonal(
                 raise ValueError("extract can be true only for Ndim >=2")
             axes = None
         else:
-            if isinstance(axis1, int) and isinstance(axis2, int):
-                if axes is not None:
-                    raise ValueError(
-                        "Either axis1/axis2 or axes must be supplied"
-                    )
-                axes = (axis1, axis2)
-            # default values for axes
-            elif (axis1 is None) and (axis2 is None) and (axes is None):
-                axes = (0, 1)
-            elif (axes is not None) and (
-                (axis1 is not None) or (axis2 is not None)
-            ):
-                raise ValueError("Either axis1/axis2 or axes must be supplied")
+            axes = (axis1, axis2)
         return self._diag_helper(offset=offset, axes=axes, extract=extract)
 
     @add_boilerplate("indices", "values")
diff --git a/cunumeric/module.py b/cunumeric/module.py
index 257fca9fa..d4a801b6a 100644
--- a/cunumeric/module.py
+++ b/cunumeric/module.py
@@ -3802,10 +3802,9 @@ def compress(
 def diagonal(
     a: ndarray,
     offset: int = 0,
-    axis1: Optional[int] = None,
-    axis2: Optional[int] = None,
+    axis1: int = 0,
+    axis2: int = 1,
     extract: bool = True,
-    axes: Optional[tuple[int, int]] = None,
 ) -> ndarray:
     """
     diagonal(a: ndarray, offset=0, axis1=None, axis2=None)
@@ -3864,9 +3863,7 @@ def diagonal(
     Multiple GPUs, Multiple CPUs
 
     """
-    return a.diagonal(
-        offset=offset, axis1=axis1, axis2=axis2, extract=extract, axes=axes
-    )
+    return a.diagonal(offset=offset, axis1=axis1, axis2=axis2, extract=extract)
 
 
 @add_boilerplate("a", "indices", "values")