Add experimental __array_module__ method

[shoyer]

xref #1565

__array_module__ (see NEP 37) is an experimental alternative to __array_function__ and __array_ufunc__ for "duck array" compatibility with NumPy that promises to be much less invasive.

Example usage, for writing a generic version of np.stack() on top of concatenate, asarray, shape attributes and array indexing:

import numpy_dispatch

def duckarray_stack(arrays):
    """This "stack" function should work with any array library, including JAX."""
    npx = numpy_dispatch.get_array_module(*arrays)  # returns jax.numpy for JAX arrays
    arrays = [npx.asarray(arr) for arr in arrays]
    shapes = {arr.shape for arr in arrays}
    if len(shapes) != 1:
        raise ValueError('all input arrays must have the same shape')
    expanded_arrays = [arr[npx.newaxis, ...] for arr in arrays]
    return npx.concatenate(expanded_arrays, axis=0)

For now, you need to use numpy_dispatch.get_array_module() to try this feature (with https://github.com/seberg/numpy-dispatch) but in the long term the hope is to merge this into NumPy proper as np.get_array_module().

My reasoning for merging it into JAX (on an experimental basis with no guarantees, of course) is that:

1. It's not invasive -- the implementation is small and self-contained.
2. No backwards compatibility issues. Unlike __array_function__ and __array_ufunc__, __array_module__ will always require an explicit opt-in by libraries that use it by calling get_array_module().
3. Other NumPy developers want evidence that this is actually feasible.
4. Scikit-Learn developers like @thomasjpfan are interested in exploring supporting scikit-learn on top of NumPy-like libraries like JAX, and experimental support for this protocol will make that easier.

Note: this PR does add numpy-dispatch as a optional testing requirement in order to verify that this works. If desired, we could remove this from CI, but installing numpy-dispatch appears to only add about 7 extra seconds of build time.

[shoyer]

I'm not quite sure it this list should include UnshapedArray or one of JAX's other abstract array types.

Their presence in the set of types checked by jax.numpy.ndarray suggests yes, but on the other hand I don't think I've ever seen them in user supplied functions transformed with JAX?

[jakevdp]

I like limiting the set for now. If it becomes clear later that UnshapedArray or others are necessary, they are easy enough to add.

[jakevdp]

Looks good! I'm in favor of provisionally adding this to help motivate this kind of feature in the broader Python array ecosystem.

On the JAX side, one detail we may need to think about is this from the Array module contract:

Unimplemented functionality should simply be omitted (e.g., accessing an unimplemented function should raise AttributeError).

This is counter to JAX's current strategy, which is to define unimplemented functions but raise NotImplementedError when they are called. I'm not sure how important that detail is. I think it's probably OK to push forward as-is and address that if it becomes an issue.

[shoyer]

Unimplemented functionality should simply be omitted (e.g., accessing an unimplemented function should raise AttributeError).

This is counter to JAX's current strategy, which is to define unimplemented functions but raise NotImplementedError when they are called. I'm not sure how important that detail is. I think it's probably OK to push forward as-is and address that if it becomes an issue.

I don't think this detail of NEP-37 (the precise error message) matters too much.

Possibly NEP-37 should be updated, since it can indeed be helpful to distinguish between a spelling error vs. an error from a function that just hasn't been implemented. Perhaps the cleanest way to do this would be module level __getattr__ from Python 3.7+ (https://www.python.org/dev/peps/pep-0562/), e.g.,

# in jax/numpy/__init__.py
def __getattr__(name):
  import numpy
  if hasattr(numpy, name):
    raise AttributeError(f'numpy.{name} has not been implemented yet in jax.numpy')
  else:
    raise AttributeError(f"module 'jax.numpy' has no attribute {name}")

[thomasjpfan]

@shoyer Thank you for working on this! I am looking forward to trying this out.