Merge pull request #21453 from jakevdp:index-update-annotation

PiperOrigin-RevId: 637938477

Author: jax authors

jax/_src/basearray.pyi

@@ -237,10 +237,10 @@ class _IndexUpdateHelper:
 
 class _IndexUpdateRef:
   def get(self, indices_are_sorted: bool = False, unique_indices: bool = False,
-          mode: Optional[str] = None, fill_value: Optional[ArrayLike] = None) -> Array: ...
+          mode: Optional[str] = None, fill_value: Optional[StaticScalar] = None) -> Array: ...
   def set(self, values: Any,
           indices_are_sorted: bool = False, unique_indices: bool = False,
-          mode: Optional[str] = None, fill_value: Optional[ArrayLike] = None) -> Array: ...
+          mode: Optional[str] = None, fill_value: Optional[StaticScalar] = None) -> Array: ...
   def add(self, values: Any, indices_are_sorted: bool = False,
           unique_indices: bool = False, mode: Optional[str] = None) -> Array: ...
   def mul(self, values: Any, indices_are_sorted: bool = False,