How to discover an array type's namespace from an input type annotation

[nicholasjng]

In machine learning workflow engines, I've come across a lot of code similar to this (with JAX as the array API implementer of choice in this example):

import jax
import jax.numpy as jnp


@task
def zeros() -> jax.Array:
    return jnp.zeros(5)


@task
def compute_sum(arr: jax.Array) -> float:
    return jnp.sum(arr).item()


@workflow
def basic_workflow():
    arr = zeros()
    compute_sum(arr)


if __name__ == "__main__":
    basic_workflow()

where the array I/O needed for the output of zeros and the input of compute_sum is abstracted away from the user in the @task decorator.

A lot of workflow libraries allow the user to customize said I/O, perhaps with an interface like the following:

class MyJAXArrayIO:
    def load(val: Any, typ: jax.Array) -> jax.Array: ...
        """Instantiate a loaded value as a JAX array."""

    def store(val: jax.Array) -> None: ...
        """Write a JAX array, e.g. to disk."""

which would be called to instantiate all of the arrays that come out of zeros and go into compute_sum, respectively. These happen to be the same in this example.

I think the array API standard would be a good fit to generalize such an I/O machinery to handle any array API standard implementing type, since it could provide a very basic load mechanism for arrays like this:

class MyArrayStandardIO(Generic[T]):
    def load(val: Any, typ: T) -> T:
        """Instantiate a loaded value as a array-API-conforming array."""
        # step 1: Look up the type's array API namespace...
        namespace = typ.__array_namespace__()
        # ... and return the value as an array in that namespace.
        return namespace.asarray(val)

Unfortunately, it seems that __array_namespace__ is an instance method only, so I cannot call it on the array type itself. Indeed, in the example of JAX, I get the following:

>>> import jax
>>> jax.Array.__array_namespace__()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/nicholasjunge/.../site-packages/jax/_src/numpy/array_methods.py", line 1133, in method
    raise NotImplementedError(f"Cannot call abstract method {name}")
NotImplementedError: Cannot call abstract method __array_namespace__

NumPy produces a different error, but I'm also unable to obtain the namespace there. As an added difficulty, JAX's array namespace is jax.numpy, but jax.Array.__module__ == 'jax'.

Is there any standards-conforming way I can deduce the array namespace from an array type annotation only?