What's your problem?

[jorenham]

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.cluster.vq

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 #22 (comment)

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.

[NeilGirdhar]

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.

Also, it would be nice to have dtype support. Jax doesn't have it, so the placeholders I'm using don't have any effect.

[jorenham]

Also, it would be nice to have dtype support.

What would that (roughly) look like in your example?

[NeilGirdhar]

What would that (roughly) look like in your example?

Ideally, just this:

ComplexArray = xpt.Array[ComplexDType]  # Or however this is done.
RealArray = xpt.Array[RealDType]

    def _r_s_mu(self) -> tuple[ComplexArray, RealArray, ComplexArray]:
        # Rest of function unchanged

And ideally, everything would be type checked.

Even better would be to make all of the classes in Efax generic on the array library somehow, and then you could have:

gamma = Gamma(jax.numpy.asarray(1.0))  # gamma: Gamma[JaxArrayNamespace]
a = gamma.pdf(b)  # Verifies that b is a jax.Array, and a has type jax.Array.

Not sure if that's possible, but that would be ideal. Maybe with Ty's intersection types, we could somehow intersect the xpt.Array[Real] with jax.Array?

[jorenham]

ComplexArray = xpt.Array[ComplexDType] # Or however this is done.
RealArray = xpt.Array[RealDType]

The # Or however this is done. is what will make up 75% of array-api-typing, in my estimation.

The big challenge is for array-api-typing to figure out what the type of ComplexDType is, and in such a way that tymypyflyright will also be able to understand it.
And the solution I had in mind for that, involves something like structural HasArrayNamespace type, and a bunch of other protocols such as HasBoolDType[_, _]. You can then obtain the actual dtype type T by matching it against HasBoolDType[HasArrayNamespace, T] in your function signature.

And ideally, everything would be type checked.

That's a hard requirement, as far as I'm concerned :)

Maybe with Ty's intersection types, we could somehow intersect the xpt.Array[Real] with jax.Array?

Intersection types will indeed be very useful when dealing with structural typing. But until the PEP is finished and accepted, we won't be able to use it, unfortunately. There has been some recent development on it: CarliJoy/intersection_examples#53.
I think we'll be able to manage without intersection types, though.

[NeilGirdhar]

You can then obtain the actual dtype type T by matching it against HasBoolDType[HasArrayNamespace, T] in your function signature.

I just hope this doesn't become a mess of templates. That would be a large cost in itself. Maybe we should sketch out what my code would look like under this system?