What's your problem?

首先查看描述中链接的 scipy-stubs issue，以及 efax/_src/distributions/complex_normal/complex_normal.py 中的 efax 示例。比较其中展示的 Array API 类型标注问题，并收集具有代表性的代码示例，以指导 array-api-typing 的改进和集成测试。当所需案例已清晰记录到足以指导类型标注工作时，即视为完成。

Let's collect a bunch of code samples from array-api libraries that array-api-typing could improve. That should help give us a better idea of what is needed. They should also be useful for (integration) testing.

Problems

scipy-stubs

See https://github.com/scipy/scipy-stubs/issues/140 for all of SciPy's functions that have (at least some) support for the Array API.

Here's an example from the scipy docs

import torch
from scipy.cluster.vq import vq

code_book = torch.tensor([[1., 1., 1.],
                          [2., 2., 2.]])
features  = torch.tensor([[1.9, 2.3, 1.7],
                          [1.5, 2.5, 2.2],
                          [0.8, 0.6, 1.7]])

code, dist = vq(features, code_book)
print(code)  # tensor([1, 1, 0], dtype=torch.int32)
print(dist)  # tensor([0.4359, 0.7348, 0.8307])

There is currently no good way to annotate this behavior in scipy-stubs. Currently, with scipy-stubs==1.15.3.0, code and dist are inferred as 1d numpy arrays of int and float dtypes, respectively. The reason for this is the torch.tensor.__array__ method, cuasing it to be interpreted as a numpy "array-like".

efax

By @NeilGirdhar from https://github.com/data-apis/array-api-typing/issues/22#issuecomment-2973309168

Consider practically any function in my efax library. For example, we have:

def _r_s_mu(self) -> tuple[JaxComplexArray, JaxRealArray, JaxComplexArray]:
    xp = array_namespace(self)
    r = -self.pseudo_precision / self.negative_precision
    s = xp.reciprocal((abs_square(r) - 1.0) * self.negative_precision)
    k = self.pseudo_precision / self.negative_precision
    l_eta = 0.5 * self.mean_times_precision / ((abs_square(k) - 1.0) * self.negative_precision)
    mu = xp.conj(l_eta) - xp.conj(self.pseudo_precision / self.negative_precision) * l_eta
    return r, s, mu

This is barely type-checked. It would be nice if xp: ArrayNamespace, r, s, k, l_eta, mu: Array, etc.