Add scalar python types to array-like types

[romanngg]

This will make isinstance called on scalars return True for both jitted and non-jitted JAX programs. Note that
isinstance(x, onp.ndarray) is False for these types, so it breaks this correspondence.

#1080

[googlebot]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it!) and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[googlebot]

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

[mattjj]

LGTM!

You could check six.PY3 and include long if we're in PY2. WDYT?

[romanngg]

Thanks Matt and James, added your suggestions!

[shoyer]

It's a little surprising that types change under jit, but it sort of makes sense if you think about it -- and I fear the fix here might be worse. What if I'm using isinstance checks to see if I can safely use NumPy attributes, like dtype and shape?

[mattjj]

@shoyer let's keep discussing this! Though we should keep in mind we're going for "best reasonable solution" rather than "perfect."

Attributes like dtype and shape work under a jit (because scalars are always 0d arrays) but not without a jit. But that doesn't seem so bad because NumPy ascribes dtypes and shapes to Python scalars already (i.e. the output of onp.result_type and onp.shape respectively). A separate argument is that using isinstance to look for attributes arguably goes against the Python duck-typing spirit.

But I am concerned if you think this 'fix' is worse. Can you say more about how this might be worse? Is it about diverging further from NumPy/Python scalar behavior, or rather about having different behavior under jit (specifically around isinstance and how it relates to attributes being available)?

Maybe another principle we should include here is the principle of least surprise. Is isinstance(3, np.ndarray) surprising?

[shoyer]

Honestly, I'm a little amazed that JAX manages to make isinstance with jax.numpy.ndarray work at all :).

Perhaps we should think about use cases? For the use case of checking for arrays/scalars vs. other non-primitive Python things inside some sort of nested Python data structure, isinstance(3, np.ndarray) seems perfectly reasonable. So maybe this is fine after all...

[shoyer]

I cannot think of circumstances when it would be useful to distinguish between Python scalars and arrays with JAX, with the possible exception of value-based casting (and that can be handled by some separate mechanism). It does seems like something like np.is_number_or_array() would be the right way to handle this, but NumPy doesn't provide that API 🤷‍♂

[hawkinsp]

Closing this PR, since it's fairly stale at this point.

I think #1080 is to some extent working as intended; jit has the effect of promoting scalars to ndarrays. So far we are mostly happy with that choice, but it is not completely transparent to introspection tools like isinstance.