Fix #1079

nv-legate · Nov 15, 2023 · 8e10a8a · 8e10a8a
1 parent 1fa075d
commit 8e10a8a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/cunumeric/array.py b/cunumeric/array.py
@@ -1675,8 +1675,6 @@ def __setitem__(self, key: Any, value: ndarray) -> None:
 
         """
         check_writeable(self)
-        if key is None:
-            raise KeyError("invalid key passed to cunumeric.ndarray")
         if value.dtype != self.dtype:
             temp = ndarray(value.shape, dtype=self.dtype, inputs=(value,))
             temp._thunk.convert(value._thunk)

diff --git a/tests/integration/test_set_item.py b/tests/integration/test_set_item.py
@@ -29,6 +29,14 @@ def test_basic():
     assert x[2] == 3
 
 
+def test_newaxis():
+    arr = num.ones((4,))
+    arr[None] = 1
+    assert np.array_equal(arr, [1, 1, 1, 1])
+    arr[None, :] = 2
+    assert np.array_equal(arr, [2, 2, 2, 2])
+
+
 ARRAYS_4_3_2_1_0 = [
     4 - num.arange(5),
     4 - np.arange(5),