removing unused variable in diagonal implementation

nv-legate · Nov 23, 2023 · 4325c55 · 4325c55
1 parent 6fda443
commit 4325c55
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 22 deletions.
diff --git 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
@@ -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")