Could we just get get_namespace typed for now?

[NeilGirdhar]

It would make a big difference just to have something like this:

Array = Any
def get_namespace(*xs: Array, api_version: str | None = None, use_compat: bool | None = None) -> types.ModuleType:

[jorenham]

Are you talking about implementing this as a runtime function, so that you could xpt.array_namespace(array)?

[NeilGirdhar]

I just meant the type annotation for array_api_compat.get_namespace. Is that something that will end up in this repository?

[jorenham]

This project will be mostly Protocols and type-guards. "Bare" annotations like these, are only useful in stub packages, which can only target a specific (runtime) package. So I don't see how that would be useful in this case.

So in other words; it's not possible to provide generic annotations for functions this way. The closest you could get is through a callable protocol, but I don't think that this would help for this get_namespace function.

[jorenham]

Did you have a specific use-case or array-api library in mind?

[NeilGirdhar]

Sorry, I don't understand where get_namespace's type annotation will ultimately be defined. Should I try to commit it to the array-api-compat project instead?

This project will be mostly Protocols and type-guards. "Bare" annotations like these, are only useful in stub packages, which can only target a specific (runtime) package. So I don't see how that would be useful in this case.

I imagined that you would ultimately do something like:

def get_namespace(...) -> Namespace: ...

class Namespace(ModuleType):
  def asarray(self, ...) -> Array: ...
  def logaddexp(self, ...) -> Array: ...
  # etc.

That way, in user code:

def softplus(x: Array, ...) -> Array:
  xp = get_namespace(x)
  return xp.logaddexp(x, ...)  # This can be type-checked because xp has type Namespace.

Did you have a specific use-case or array-api library in mind?

Most functions that use the Array API will call get_namespace. It's useful if its return type is not Unknown. ModuleType would be an improvement. And the above Namespace would be ideal?

[jorenham]

Should I try to commit it to the array-api-compat project instead?

Ah I see; the issue is that array-api-compat is untyped. That's not something that array-api-typing could help with unfortunately. There are two other options though:

• contribute annotations to array-api-compat and add a py.typed. This could either be inline annotations (this is the preferred approach, as that'll allow runtime type-checkers to understand it as well), or as .pyi "stubs", which is what numpy is currently doing. Note the py.typed should could contain the contents partial iff not everything is annotated. See https://typing.readthedocs.io/en/latest/spec/distributing.html#packaging-typed-libraries for details.
• Create a separate stub package, array-api-compat-stubs. This makes the type annotations opt-in, and allows you to independently publish releases.

[NeilGirdhar]

Right. I honestly thought that's what this project would be providing. What kind of things are you planning on annotating here? And why wouldn't you also want to annotate the get_namespace method?

Also, where will the Namespace be annotated? Here, or somewhere else?

[jorenham]

I imagined that you would ultimately do something like:

def get_namespace(...) -> Namespace: ...

class Namespace(ModuleType):
  def asarray(self, ...) -> Array: ...
  def logaddexp(self, ...) -> Array: ...
  # etc.

That way, in user code:

def softplus(x: Array, ...) -> Array:
  xp = get_namespace(x)
  return xp.logaddexp(x, ...)  # This can be type-checked because xp has type Namespace.

Protocols like Namespace and Array are indeed what we intend do develop here in array-api-typing. But until we have those, the "gradual typing" approach would be to use either Any, or in case of the Namespace, use types.ModuleType (whose attributes are all Any).

[jorenham]

Right. I honestly thought that's what this project would be providing.

I wish we could, but it's not possible within the Python typing system: You can only annotate functions from either dedicated stub packages, or from within the package itself.

[NeilGirdhar]

Protocols like Namespace and Array are indeed what we intend do develop here i

Got it. Perfect, that's what I figured.

I wish we could, but it's not possible within the Python typing system.

Got it, that makes sense.

So, it's probably unreasonable for me to ask the array-api-compat project to depend on this one to provide the annotation that says get_namespace returns array_api_typing.Namespace, right? So we would need a whole stub project just for that single annotation?

Just want to make sure I understand.

[jorenham]

o we would need a whole stub project just for that single annotation?

Eventually; yes. If type-checkers can't understand it, then array-api-typing also won't be of any help. So it's only typed libraries that array-api-typing will be able to help with.

But for now, you could gradually introduce annotations, starting with e.g. the get_namespace function. Just don't forget to add a py.typed with partial\n as contents. See https://typing.readthedocs.io/en/latest/spec/distributing.html#partial-stub-packages.

[jorenham]

Anyway, let me know if there's anything with array-api-compat typing stuff I can help out with :)

[NeilGirdhar]

Perfect, thanks for explaining. I think we won't need partial since there are only a few functions exported by that project and we can probably just type them all.

Just confirming: you think it would be unreasonable for array-api-compat to depend on array-api-typing? Wouldn't all users benefit from the annotations? Looks like they've started with some rudimentary protocols over there in array_api_compat/common/_typing.py

Also, is it possible to do something like:

if TYPE_CHECKING:
  from jax import Array as JaxArray  # Somehow this should not cause any type errors even if Jax isn't installed.
  def is_jax_array(x: Array) -> IsType[JaxArray]: ...
else:
  def is_jax_array(...) -> ...: ...

[jorenham]

you think it would be unreasonable for array-api-compat to depend on array-api-typing?

I can imagine that doing so would allow for writing narrower annotations, so it might indeed be a good idea.