Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

Type annotations not working for DeviceArray

See original GitHub issue

MyPy says that jax.interpreters.xla.DeviceArray has no attribute “shape”, is not indexable, etc. It seems that it points to DeviceArrayBase—not DeviceArray, which has the annotations.

I know that type annotations are a work-in-progress. This is just a placeholder issue 😃

Issue Analytics

State:
Created 2 years ago
Comments:8 (4 by maintainers)

Top GitHub Comments

1reaction

NeilGirdharcommented, May 14, 2021

Okay, that makes sense. Thanks for explaining it. I’ll just block the MyPy errors!

0reactions

jakevdpcommented, May 13, 2021

I think there’s a distinction here between “DeviceArray as xla operand” and “DeviceArray as duck type for numpy arrays”. The former should not have numpy-like attributes, the latter should.

This might argue that we need another base class in the hierarchy that defines the interface to the two DeviceArray implementations; I think the reason we haven’t done that is that the two implementations are temporary; somewhere there’s a TODO to unify them again.

Would it serve your purposes to locally disable mypy type checking, and only use the DeviceArray annotation for dispatch?