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.