y is always defined

Author: crusaderky

src/array_api_extra/_funcs.py

@@ -759,7 +759,7 @@ def get(
     def _update_common(
         self,
         at_op: str,
-        y: Array = _undef,
+        y: Array,
         /,
         copy: bool | None = True,
         xp: ModuleType | None = None,
@@ -783,8 +783,8 @@ def _update_common(
         if copy:
             if is_jax_array(x):
                 # Use JAX's at[] or other library that with the same duck-type API
-                func = getattr(x.at[self._idx], at_op)
-                return func(y) if y is not _undef else func(), None
+                func = getattr(x.at[idx], at_op)
+                return func(y), None
             # Emulate at[] behaviour for non-JAX arrays
             # with a copy followed by an update
             if xp is None: