Reuse alternate index syntax

Author: crusaderky

src/array_api_extra/_funcs.py

@@ -672,8 +672,7 @@ def __getitem__(self, idx: Index, /) -> "at":  # numpydoc ignore=PR01,RT01
         if self._idx is not _undef:
             msg = "Index has already been set"
             raise ValueError(msg)
-        self._idx = idx
-        return self
+        return at(self._x, idx)
 
     def _update_common(
         self,

tests/test_at.py

@@ -127,6 +127,11 @@ def test_alternate_index_syntax():
     a = np.asarray([1, 2, 3])
     assert_array_equal(at(a, 0).set(4), [4, 2, 3])
     assert_array_equal(at(a)[0].set(4), [4, 2, 3])
+
+    a_at = at(a)
+    assert_array_equal(a_at[0].add(1), [2, 2, 3])
+    assert_array_equal(a_at[1].add(2), [1, 4, 3])
+
     with pytest.raises(ValueError, match="Index"):
         at(a).set(4)
     with pytest.raises(ValueError, match="Index"):