Fix annotations

[NeilGirdhar]

This pull request probably needs some reviewer guidance. It fixes all (but one) of the MyPy errors after the jit-annotation PR. I will comment on various changes in line.

The final type error should be fixed by this PR.

[jakevdp]

Overall comment here: the approach in this PR is one I've tried in other contexts and often ended up having to roll back: namely, you're trying to fix type check errors by changing the runtime behavior of the function. It's very likely that such changes will have unexpected effects in unanticipated corners of downstream JAX use, and so this kind of PR will likely lead to a rollback.

Instead, we should fix type checking errors by fixing type annotations, without changing the runtime logic. That's a much safer change in general.

[NeilGirdhar]

Instead, we should fix type checking errors by fixing type annotations, without changing the runtime logic. That's a much safer change in general.

Okay, I know you're right. Thanks for being patient with this PR. I'll scan over this change and try to eliminate behavior changes.

[jakevdp]

Also, I think some of this PR does not make sense separate from the jit annotation change, while some of it does make sense to be separate. A lot of the little fiddly changes should probably be landed in concert with the JIT change. It's not worth doing those separately now, because the jit change will probably evolve before we land it, so a different set of fiddly changes will eventually be needed.

Some changes here make sense as stand-alone though: for example adding real type annotations to properties, or adding overloads for atleast_1d, etc. Each of those should be done as a separate PR, mainly because the more changes you lump into a single PR, the higher probability you have of having to roll back the whole thing because one part broke something downstream.

[NeilGirdhar]

Makes perfect sense.

Just to let you know that I have tested this PR by rebasing it onto the jit annotation change, making sure it still has no type errors, and then rebasing back onto main.

Do you want to let me know how you'd like this PR split?

[NeilGirdhar]

@jakevdp FYI this change is on top of #18465, so I made some of your changes there and rebased over it.

[jakevdp]

FYI this change is on top of #18465, so I made some of your changes there and rebased over it.

Sorry, I should have realized that.

[NeilGirdhar]

@jakevdp Just checking that we're on the same page. These remaining changes are just what you described as the "fiddly" changes for the annotate Jit PR that you said would eventually change when we go to do that. Sorry for the rebase, I just wanted to make sure that it was clean with Ruff.

Let me know if there's anything you want me to split off into another PR though, and I'm happy to do that.

[jakevdp]

Thanks! I think all this looks pretty reasonable now. We'll need to do from __future__ import annotations in a few places to fix the Python 3.9 errors, and there are a few existing comments to address, but other than that I'm happy to merge this. Thanks so much for working on this!

[NeilGirdhar]

@jakevdp Is this okay to merge?

[jakevdp]

Let's see what the tests say 😀

[jakevdp]

Test failures are real, it looks like the out_indices change is modifying the logic.