Skip to content

What's your problem? #22

Open
Task
Open
What's your problem?#22
Task
Labels
help wantedExtra attention is needed✅ testsAdd, update, or pass tests.

Description

@jorenham
jorenham
opened on Jun 13, 2025

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 scipy/scipy-stubs#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 #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.

Activity

  1. jorenham
    added the
    Task
    issue type on Jun 13, 2025
  2. jorenham
    added
    help wantedExtra attention is needed
    ✅ testsAdd, update, or pass tests.
    on Jun 13, 2025
  3. jorenham
    pinned this issue on Jun 13, 2025

  4. NeilGirdhar commented on Jun 14, 2025

    @NeilGirdhar
    NeilGirdhar
    on Jun 14, 2025
    Contributor

    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.

  5. jorenham commented on Jun 15, 2025

    @jorenham
    jorenham
    on Jun 15, 2025
    MemberAuthor

    Also, it would be nice to have dtype support.

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

  6. NeilGirdhar commented on Jun 15, 2025

    @NeilGirdhar
    NeilGirdhar
    on Jun 15, 2025
    Contributor

    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?

  7. jorenham commented on Jun 15, 2025

    @jorenham
    jorenham
    on Jun 15, 2025
    MemberAuthor

    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.

  8. NeilGirdhar commented on Jun 15, 2025

    @NeilGirdhar
    NeilGirdhar
    on Jun 15, 2025
    Contributor

    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?

  9. nstarman commented on Aug 22, 2025

    @nstarman
    nstarman
    on Aug 22, 2025
    Collaborator
    1. Have method-by-method Can/Has{X} Protocols for the methods and attributes of the array object in the Array API. I want granular descriptions of what my functions need, e.g. (very loosely)
    def square_minus_one(x: ArrayMul & ArraySub, /):  xp.square(x) - 1
    1. Have useful intersection types like Array and partially filled in types that describe numerical versus boolean arrays.
    def square_minus_one(x: Array, /) -> Array:  xp.square(x) - 1
    1. Shape typing would be great.
    def square_minus_one(x: Array[Shape], /) -> Array[Shape]:  xp.square(x) - 1
    1. Have function-by-function Does{X} Protocols for the functions in the Array API.
    def apply_fn(func: DoesAsType, x):  func(x)
    1. Have useful intersection types like ArrayNamespace.
    2. Have Array and ArrayNamespace work together 🤝 .

  10. lucascolley commented on Aug 22, 2025

    @lucascolley
    lucascolley
    on Aug 22, 2025
    Member

    I've had multiple different people asking about shape-typing at EuroSciPy btw :)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    help wantedExtra attention is needed✅ testsAdd, update, or pass tests.

    Type

    Task

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions