Array API compatibility

[yaugenst-flex]

This is more of a question rather than an issue, but is there any plan or opinion on supporting Array API compatibility in autograd? Specifically, I'm wondering about the possibility of implementing methods like __array_function__ and __array_ufunc__ for ArrayBox, in a manner similar to the __array_namespace__ discussed in PR #647.

Adding such compatibility would enable autograd to seamlessly interoperate with libraries that follow the Array API standard, such as xarray, and potentially allow differentiating through operations involving those libraries. This would mean that automatic differentiation could potentially "just work" for any library that dispatches to the underlying array implementation.

[fjosw]

Hi @yaugenst-flex, we don't have any plans (and likely not the resources) to work on Array API compatibility but we would be happy about all contributions that extend or improve autograd's usability.

[agriyakhetarpal]

Hi @yaugenst-flex, thanks for this request! Personally, I am quite interested in the array API standard, now that several array libraries are working towards supporting the 2023 edition, and I would be keen to explore array API compatibility with Autograd. I would be keen to contribute towards preliminary support and testing infrastructure for it myself, or help with code reviews for contributions, which are always welcome – as Fabian has mentioned. That said, I don't know anything beyond the basics of the standard, so I think it would be nice to seek a plan of what needs to be done after #647 and establish conformance, and a sense of what changes Autograd in its current state needs. We could then consider this in a v2 release for Autograd, or perhaps earlier, via a v1.8.0 release.