Better Numpy API (interoperability between ML frameworks)
Labels: triaged, module: numpy

## Context

For better interoperability between ML frameworks, it would be great if `torch` API matched more closely `numpy` API (like `tf.experimental.numpy` and `jax.numpy`).

This is a highly requested features, like: #2228 (~100 upvotes), #50344, #38349,... Those issues have since been closed even though there's still many issues remaining.

This is even more relevant with the numpy API standard ([NEP 47](https://numpy.org/neps/nep-0047-array-api-standard.html)): The goal is to write functions once, then reuse them across frameworks:

```python
def somefunc(x, y):
    xnp = get_numpy_module(x, y)  # Returns `np`, `jnp`, `tf.numpy`, `torch`
    out = xnp.mean(x, axis=0) + 2*xnp.std(y, axis=0)
    return out
```

Our team has multiple universal libraries which support both `numpy`, `jax` and `TF` (like [`dataclass_array`](https://github.com/google-research/dataclass_array) or [`visu3d`](https://github.com/google-research/visu3d)).

We've been experimenting adding `torch` support recently but encountered quite a lot of issues (while `tf.numpy`, `jax` worked (mostly) out of the box). Here are all the issues we're encountered:

## numpy API issues

Some common methods (present in `np`, `jnp`, `tf.numpy`) are missing from `torch`:

- [ ] `torch.array`: like `x = xnp.array([1, 2, 3])`) (alias of `torch.tensor`)
- [ ] `torch.ndarray`: like `isinstance(x, xnp.ndarray)` (alias of `torch.Tensor`)
- [ ] `Tensor.astype`: like `x = x.astype(np.uint8)`
- [ ] `torch.append` (https://github.com/pytorch/pytorch/issues/64359): https://numpy.org/doc/stable/reference/generated/numpy.append.html
- [ ] `torch.expand_dims`: https://numpy.org/doc/stable/reference/generated/numpy.expand_dims.html (alias of `torch.unsqueeze`)
- [ ] `torch.around`: https://numpy.org/doc/stable/reference/generated/numpy.around.html
- [ ] `torch.concatenate`: https://numpy.org/doc/stable/reference/generated/numpy.concatenate.html

Behavior:

- [ ] torch should accept `np.dtype` everywhere `torch.dtype` is valid (e.g. `torch.ones((), dtype=np.int32)`) (See: #40568).
- [ ] `torch.dtype` should be comparable with `np.dtype`: `tf.int32 == np.int32` but `torch.int32 != np.int32`. This allow to have agnostic comparison: `x.dtype == np.uint8` working for all frameworks.
- [ ] `x[y]` fail for `y.dtype == int32` (raise `IndexError: tensors used as indices must be long, byte or bool tensors`) but works in TF, Jax, Numpy (https://github.com/pytorch/pytorch/issues/83702).
- [ ] `axis=()` (like in `x.mean(axis=())`) return a scalar for `torch` but is a no-op for `np`, `jnp`, `tf` (besides returning float for int array, see bellow). (https://github.com/pytorch/pytorch/issues/29137)
- [ ] `torch.ones(shape=())` fail (expect `size=`), but `xnp.ones(shape=())` works (in other frameworks). Same for `torch.zeros`,...

Casting:

- [ ] `x.mean()` currently require explicit dtype when `x.dtype == int32` (`x.mean(float32)`). Other frameworks default to the default float type (`float32`)
- [ ] `torch.allclose(x, y)` currently fail if `x.dtype != y.dtype`, which is inconsistent with `np`, `jnp`, `tf` (this is very convenient in tests `np.allclose(x, [1, 2, 3])`)

Other differences (but not critical to fix):

- [ ] Mixing `torch` and `np.array` fail (#46829). Both TF and Jax support `tf.Tensor + np.array`.
- [ ] Mixing float and double
- [ ] Would be nice if `torch.asarray` was supporting `jax`, `tf` tensors (and vice-versa).

## Testing and experimenting

Those issues have been found in real production code. Our projects have a `@enp.testing.parametrize_xnp` to run the same unittests on `tf`, `jax`, `numpy`, `torch` to make sure our code works on all backends...

For example: https://github.com/google-research/visu3d/blob/89d2a6c9cb3dee1d63a2f5a8416272beb266510d/visu3d/math/rotation_utils_test.py#L29

In order to make our tests pass with `torch`, we had to mock `torch` to fix some of those behaviors: https://github.com/google/etils/blob/main/etils/enp/torch_mock.py

Having a universal standard API that all ML frameworks apply would be a great step towards. I hope this issue is a small step to help toward this goal.


cc @mruberry @rgommers