Add stubs for major protocols #20
Conversation
There was a problem hiding this comment.
I think this is a great start!
Thanks @NeilGirdhar.
Maybe we can get @34j to comment on how automatic generation might integrate with this project. I imagine that filling out ArrayNamespace would be the prime target.
100% agree. I thought about it, but I wasn't sure how it would work.
That makes sense! |
|
@NeilGirdhar can we wait to merge this in until #18 is in? Hopefully we can get some CI up and running to test these additions :). |
|
Sounds great, thanks for the thorough review! I know that @jorenham wants to review this as well 😄 |
Yea that's right, thanks for the ping. I'll have a look at it soon. |
There was a problem hiding this comment.
Thanks for this PR!
This PR looks pretty similar to what I tried to build before in optype: https://github.com/jorenham/optype/tree/array-api/optype/array_api. There, I also started with the same enthusiasm as I see here. After spending a good amount of time on it, I realized that none of the protocols would be able to solve any problem in practice. So I'm a bit worried that this will happen again.
One of the mistakes I made, was that I started without being able to "see" anything: I should have instead gone with a test-driven approach. Because without integration tests you can't tell whether the relevant libraries will be compatible with the protocols. And in this optype case, it later realized none of the relevant libraries were compatible.
The other big mistake I made, is that I started without knowing the exact problem that I was trying to solve. This is something I still tend to underestimate, so I know how difficult it is.
And as hypocritical as it might sound, that's why I think we should first figure out what exactly the problems are that we're trying to solve with array-api-typing.
src/array_api_typing/_device.py
Outdated
| @final | ||
| class Device(Protocol): | ||
| pass |
There was a problem hiding this comment.
I agree it would be useful to be able to structurally type a device, but I don't think that that's possible given the current spec. Right now, this protocol is equivalent to Device = object. And since I don't see any way how we could expand it later on, I think it might be better to leave the the Device out. Because I'm worried that this object alias will give users a false sense of type-safety.
There was a problem hiding this comment.
I see your point about type-safety illusions, but I think we should at least have Device = object so that other libraries have something to point to. Even without type safety, just writing xpt.Device is a nicer annotation.
There was a problem hiding this comment.
This is what I do in my code when I need to annotate a Jax random number generator key. I find that key: KeyArrayis a lot easier to read that key: Array. And, in theory, if we ever get dtype support for Jax array annotations, then we make KeyArray = jax.Array[KeyDType].
src/array_api_typing/_dtype.py
Outdated
There was a problem hiding this comment.
Maybe I'm wrong here, but I think that the spec for dtypes is too broad for us to be able to structurally describe it. So my previous comment for Device also applies here.
There was a problem hiding this comment.
Right, let's discuss this along with Device.
src/array_api_typing/_namespace.py
Outdated
|
|
||
| def asarray( | ||
| self, | ||
| obj: Array | complex | NestedSequence[complex] | SupportsBufferProtocol, |
There was a problem hiding this comment.
This is a pretty heavy restriction. Because this will now reject your asarray method unless Array | complex | NestedSequence[complex] | Buffer is assignable to the type of your obj. That's because function domain types are contravariant, which is a consequence of the Liskov Substitution Principle (LSP).
There was a problem hiding this comment.
Right, that's the promise that's shown on the Array API documentation. Did you want a narrower interface?
src/array_api_typing/_namespace.py
Outdated
| def astype( | ||
| self, | ||
| x: A, | ||
| dtype: DType, | ||
| /, | ||
| *, | ||
| copy: bool = True, | ||
| device: Device | None = None, | ||
| ) -> A: ... |
There was a problem hiding this comment.
Even if dtype is passed, the array type A will be unaffected. That, when combined with the invariance of A, means that we will not (and cannot) support array types with generic dtype parameter. I'm not sure if that's what we want.
There was a problem hiding this comment.
That's a really good point. What I hoped for was that:
x: jax.Array
xp = array_api_compat.array_namespace(x)
y = xp.asarray(0.0)
reveal_type(y) # jax.ArrayIs there a way to have that without also promising that the dtype of x and y are the same?
I guess you could regenerate the entire set of typing stubs for each backend. Is that what we need to do?
There was a problem hiding this comment.
Why did you choose to make this a package instead of a module?
There was a problem hiding this comment.
I wasn't sure if we would end up with more modules. But you're right, YAGNI, so should I switch it to a module?
Yes, this is one of the main issues with the design decision to design the libraries first and the annotations second. If we had done it the other way around with the interfaces first and then have the libraries explicitly inherit from the interfaces, there wouldn't be a problem. However, I think we already have access to an easy solution. We don't need the libraries to conform to the protocols. We just need the interface returned by From there, we can apply soft pressure on the libraries to conform more so that the shims slowly disappear. What do you think?
Yes, that makes sense. We're not writing code for its own sake. For me, I would like my efax library to become typed. Currently, I've typed the arrays as Also, I think it means that This is why I simply lifted the annotations from What do you think? |
|
Sorry @NeilGirdhar. Rebase required. |
|
Okay, I'm sorry, but I have to make a new PR. I'll link it here |
No description provided.