Add __array_namespace__ to ArrayBox

[yaugenst-flex]

Small change that adds the __array_namespace__ method from the Array API standard to autograd's ArrayBox.

This makes autograd (partially) compatible with libraries that check for array API compatibility, such as xarray, because xarray will coerce types that do not implement the array API into regular numpy arrays, see relevant discussion here.

The following code snippet:

import xarray as xr

import autograd.numpy as np
from autograd import grad


def f(x):
    data = xr.DataArray(x, dims=range(x.ndim)).data
    print(data)
    return np.sum(data)


x = np.random.uniform(-1, 1, 10)
grad(f)(x)

prints

[<autograd.numpy.numpy_boxes.ArrayBox object at 0x7f9465347640>
 <autograd.numpy.numpy_boxes.ArrayBox object at 0x7f94652f7000>
 <autograd.numpy.numpy_boxes.ArrayBox object at 0x7f94652f7c40>]

currently, because the traced ArrayBox is being coerced into a regular numpy array of dtype=object with ArrayBox elements. Differentiating through this works (with some limitations), but is very inefficient.

With the addition suggested in this PR, the snippet instead prints:

Autograd ArrayBox with value [1. 1. 1.]

i.e. data is now a "proper" ArrayBox.

I might also be interested in working on a full xarray wrapper in the future, but for now just this change would already help a lot.

[fjosw]

Hi @yaugenst-flex, thanks for contributing to autograd! I made the type hints compatible with python<3.10 and now all tests seem to pass 👍 Any objections against merging this @agriyakhetarpal @j-towns ?

[j-towns]

No objections from me.

[agriyakhetarpal]

No objections from me either! Thanks, @yaugenst-flex!

[agriyakhetarpal]

I've created https://github.com/HIPS/autograd/milestone/1 to track things we can add to a v1.8.0 release (or to see which things we can split up and add to a v1.7.1 release). This sounds like a useful improvement.

[yaugenst-flex]

This is great, thanks!