Type hint `at` property in `basearray.pyi`

[soraros]

Type hint jax.Array.at to aid IDEs.

Didn't submit a PR directly because

1. Not sure if there is any circular import issue
2. _IndexUpdateHelper is in internal, do we change the name or __name__ to make it look better?

General question for JAX developers: is there a 'canonical' name for the at-properties? jax.Array.at is not convenient for verbal communication. Maybe, functional array index accessor? I really want a name that is searchable, and I could refer to, say when answering questions in Github discussions.

[jakevdp]

Thanks for raising this! I think I left this out because circular imports make it tricky, but we should be able to figure it out.

The name we use now is "index update helper", but that's a bit of a mouthful. If you search for "jnp.ndarray.at" you find this page: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html

We'd be open to suggestions for making that more discoverable.

[soraros]

The name we use now is "index update helper"

I don't think it's a bad name. As for discoverability, I'd suggest inserting the exact phrase in ndarray.at document, Array class, jax.ops document, and maybe also in the shape bits. Something like

Index update helper
 or the `.at`-property

in CS paper title fashion.

I was trying to sell the array-api people this api, now I can tell them the name of the product.

[jakevdp]

Unfortunately, the header on the ndarray.at doc page is generated by sphinx autodoc, so it's not easy to change it. We can certainly add more words in the dosctring though.

[soraros]

We can certainly add more words in the dosctring though.

Yes, this is what I meant.