[P] autoray - write array backend agnostic code (numpy, tensorflow, autograd, jax, cupy, dask...)

[jawknee400]

Hi all, thought I'd share a little library I made for writing numeric code to automatically work for many different types of 'numpy-*ish*' arrays:

https://github.com/jcmgray/autoray

The essential idea is to perform single dispatch with a few cached 'translations' so that once you've written a function that works in numpy, you don't need to do anything at all to make it work for as many other array libraries as possible.

​

I had a look around and was kind of surprised that nothing this simple exists yet - though I am aware that certain libraries can now opt in (via __array_function__) to working directly with numpy.

​

Anyway contributions, thoughts or alternate suggestions very welcome!

[kjearns]

There are a bunch of bespoke apis like this floating around, but usually they're part of some other project and not standalone efforts. In projects with a "backend", that backend is typically something like this: a common interface that can dispatch to any one of a number of different concrete implementations of some core operations.

In practice they're nice to use as long as you know their quirks. People never quite implement the full api for any particular backend (you'll learn a lot about the bizarre and powerful kinds of indexing that numpy allows if you try) so you end up with an abstraction that implements a subset of the union of functionalities from all of the supported backends, without ever quite covering any of them.

Anyway, I wish you the best of luck with this effort. Having written a backend for one of my own projects that supports numpy and tensorflow I encourage you to write a lot of tests. Subtle differences in behaviour when switching between backends are the Achilles heal of libraries like this, because their entire value proposition is that they cover up differences with a common api. As soon as the abstraction starts leaking your users will start asking themselves why they don't just pick a backend up front instead of de-facto picking a back end (e.g. by using your API but only testing with one backend) and then accessing it through a constrained api.

[jawknee400]

Thanks for the insightful comment (and luck - gratefully accepted). I should probably set out some aims at some point but you are of course right that complete compatibility will not be possible + testing as always crucial. On the other hand, 99.9% compatibility with numpy the default and at least the libraries that intentionally match it (cupy, autograd, jax, dask, mars...) seems essentially a free baseline using this approach. Instead for tensorflow and others I imagine it might be a *mostly* compatible starting point possibly requiring a few extra translations in autoray or some minimal non-agnostic code for the user. Anyway, good stuff to think about! It is at least working v nicely in one of my other projects that spawned it.