使用 numpy 和 jax 进行非传递子类化

Question

我的问题很简单：

>>> isinstance(x, jax.numpy.ndarray)
True
>>> issubclass(jax.numpy.ndarray, numpy.ndarray)
True
>>> isinstance(x, numpy.ndarray)
False

?

现在我就扯远了，让SE接受我的合理问题。

Answer 1

出现这种情况的原因是 jax.numpy.ndarray 使用元类覆盖实例检查：

class _ArrayMeta(type(np.ndarray)):  # type: ignore
  """Metaclass for overriding ndarray isinstance checks."""

  def __instancecheck__(self, instance):
    try:
      return isinstance(instance.aval, _arraylike_types)
    except AttributeError:
      return isinstance(instance, _arraylike_types)

class ndarray(np.ndarray, metaclass=_ArrayMeta):
  dtype: np.dtype
  shape: Tuple[int, ...]
  size: int

  def __init__(shape, dtype=None, buffer=None, offset=0, strides=None,
               order=None):
    raise TypeError("jax.numpy.ndarray() should not be instantiated explicitly."
                    " Use jax.numpy.array, or jax.numpy.zeros instead.")

(view source)

你的代码 returns 之所以这样做是因为你有一个 x 值，它不是 numpy.ndarray 的实例，但是这个 __instancecheck__ 方法returns 正确。

为什么在 JAX 中使用这种诡计？好吧，出于 JIT 编译、auto-differentiation 和其他转换的目的，JAX 使用 stand-in 称为 tracers 的对象，这些对象看起来和行为都像一个数组，尽管实际上不是一个数组。这种对实例检查的覆盖是 JAX 用来进行此类跟踪的技巧之一。

使用 numpy 和 jax 进行非传递子类化

Nontransitive subclassing with numpy and jax

python

numpy

jax