I implemented custom dispatchers as discussed in #3.

I wrote some tests to confirm it works as expected. One potential caveat, functions like np.stack can also eat a one-dimensional array, in which case it acts like the identity function. In those cases backend is inferred from first element of this one-dimensional array by join_array_dispatcher. This works correctly for all supported backends except sparse, where backend is inferred as numpy. This is however not the intended usage of stack anyway, so I don't think this is a problem.

For join_array_dispatcher we could optionally check if args[0][0] even exists in the first place, and if not just return args[0]. This is only relevant if you want to supply it empty sequences or 0-dimensional arrays (which numpy doesn't allow, for example). If so, one can check this with hasattr(args[0], '__getitem__'). This is a bit nicer than a try: / except: block, especially since we have to catch both IndexErrors and TypeErrors.

It seems my auto formatter black was wreaking havoc on your code. Biggest thing it did is replacing single quotes for double quotes everywhere. If this bothers you, I will do a commit with only my actual changes. Otherwise, do you use a particular auto formatter for this project?

I also took the liberty of adding translations for np.take. I'm a bit confused about the syntax of make_translator. For torch implementation of take (i.e. index_select) the names of arguments and order of arguments is both different, and I'm not 100% sure I did it correctly (at least it seems to work).

This adds two nice and fairly natural features to autoray, any feedback on interface welcome!

Lazy Computation (LazyArray)

from autoray import lazy
... 
lx = lazy.array(x)
lf = fn(lx)
lf.compute()

If you write a function / algorithm with do calls, then you can trace through the entire thing lazily and:

• [x] e.g. check max array size encountered, number of calls to specific functions
• [x] plot the computational graph
• [x] identify and cache shared intermediates only (with lazy.shared_intermediates(): ...)
• [x] perform all constant parts of a computational graph

This is all implemented in a very lightweight manner making it about 100x faster than e.g. dask (which of course offers many other features) on some examples, and suitable for computational graphs with >100,000s nodes.

TODO:

• [ ] implement a few remaining functions
• [ ] eventually might be possible to estimate FLOPs etc
• [ ] eventually might be useful to execute the computational graph with e.g. ray

Auto compilation / unified JIT interface (@autocompile)

@autocompile
def fn(x, y):
    ...

fn(x, y, backend='torch')

The aim here is to have a single decorator for marking an autoray function to be JIT compiled, which makes it v easy to switch backends ('jax', 'tensorflow', 'torch') or just dispatches to the correct one automatically. cc @RikVoorhaar and #5.

• [x] currently the signature must be fn(*arrays) -> array
• [x] there is a 'python' backend that 'simply' uses the LazyArray computational graph to compile and exec an unravelled version of the function, with shared intermediates, folded constants + surrounding logic stripped away

TODO:

• [ ] CompilePython probably needs to dispatch based on the input array shapes (overhead is very minimal)

fixes #7

This seems to work. Only thing I'm not satisfied with is the comparison

A = fn(*args, **kwargs)
if (dtype is not None) and (dtype != standard_type):
    ...

For example if standard_dtype='float64' and dtype=np.float64 this will still trigger, but it shouldn't. I have three ways around this:

• Store standard_type not as string but as backend specific dtype object. Then compare standard_typetoA.dtype` instead.
• Compare standard_type to get_dtype_name(A).
• Make a cached function for dtype comparisons across backends

As mentioned in #5

• Translating np.diff didn't seem worthwhile. There is tf.experimental.numpy.diff, but it's only on newest version of tensorflow, and probably in a later version it will become just tf.diff.
• Torch and tensorflow wrappers for split can probably be unified into one function, but this would mean wrapping a translation around the wrapper I made for torch since syntax is slightly different. Both typically take a list as input if using sections to split the array, so I don't think my wrapper brings that much extra overhead.
• For torch I didn't use tensor_split since it's not yet in the stable release. Maybe one can do different things in autoray depending on torch version number, but that doesn't seem very elegant.
• I couldn't run any of the CuPy tests on this machine since it doesn't have a GPU

When using concatenate, the result is always converted to numpy arrays. E.g.

A = np.random.normal(size=(10,10),like='tensorflow')
B = np.random.normal(size=(10,10),like='tensorflow')
concat = ar.do('concatenate',(A,B),axis=0)
type(concat)
>> numpy.ndarray

This can be mitigated by instead doing

ar.do('concatenate',(A,B),axis=0,like=ar.infer_backend(A))

but this is a bit unwieldy. The problem is that the argument (A,B) is a tuple, which belongs to backend builtins, which in turn always gets inferred as numpy by infer_backend.

This problem applies to any function whose first argument is a list/tuple of arrays. I know at least that this applied to concatenate, einsum and stack. For einsum I just opted to call opt_einsum directly, which does correctly infer backend in this case, but that is besides the point.

I can see several possible approaches:

1. Make an exception in the way backend is inferred for these specific functions. When using ar.register_function the user should also be able to indicate the function is of this type.
2. In _infer_class_backend_cached make a specific check for builtins: we check if the item is iterable, if so we check the backend of the first element. If it is again builtins, then leave it as is, but if it is something else then return that backend instead.
3. Do nothing, but explicitly mention this behavior in the README.

I'm partial to the second option, as I don't expect it to have too many side-effects. If you want I can do a PR.

Problem

The api for the numpy.ndarray transpose attribute allows it to permute an arbitrary number of indices into an arbitrary order. However, the torch.Tensor transpose attribute assumes a matrix and therefore only accepts two indices. This means something like the following will fail:

import numpy
import torch
from autoray import do, transpose

Ttorch = torch.zeros([2,3,4,5])
Tnp = numpy.zeros([2,3,4,5])

print(Tnp.transpose([2,1,3,0]).shape)   # gives (4,3,5,2), as expected
print(transpose(Tnp, [2,1,3,0]).shape)  # also gives (4,3,5,2)
print(Ttorch.transpose([2,1,3,0]).size()) # this fails with a TypeError
print(transpose(Ttorch, [2,1,3,0]).size())  # which means this also fails

Solution

The correct torch.Tensor attribute is permute, which has the same exact behavior as numpy.ndarray.transpose. This means that something like the following will do what we want:

import numpy
import torch
from autoray import do, transpose

Ttorch = torch.zeros([2,3,4,5])
Tnp = numpy.zeros([2,3,4,5])

print(Tnp.transpose([2,1,3,0]).shape)   # gives (4,3,5,2), as expected
print(transpose(Tnp, [2,1,3,0]).shape)  # also gives (4,3,5,2)
print(Ttorch.permute(2,1,3,0).size())  # also gives (4,3,5,2)

Proposed code change

I'm not sure that there is a way to incorporate this behavior in a clean, non-invasive manner. As far as I understand, the _module_aliases and _func_aliases dictionaries are not applicable since permute is only an attribute of torch.Tensor (i.e. there is no torch.permute(torch.Tensor, *args)). This therefore seems to necessitate direct modification of the autoray.transpose function (line 308). The following patch works, but it's not very clean:

current code:

def transpose(x, *args):
    try:
        return x.transpose(*args)
    except AttributeError:
        return do('transpose', x, *args)

patched code:

def transpose(x, *args):
    backend = infer_backend(x)
    if backend == 'torch':
        return x.permute(*args)
    else:
        try:
            return x.transpose(*args)
        except AttributeError:
            return do('transpose', x, *args)

The inherent challenge is that we need to alias x.transpose() to x.permute() when x is a torch.Tensor. If there is a better way than what I have suggested, let me know!

(p.s.) I found this problem via an error I obtained in quimb. I was trying to fuse multiple bonds of a quimb Tensor when using pyTorch as the backend, and this problem arose.

@jcmgray Here is my attempt at implementing sparse support in autoray. Let me know if something needs to be changed significantly -- happy to help make this work.

Something I ran into is that different backends prefer either single or double precision. I personally need double precision, or at least prefer to consistently use one precision. This is also much more fair for benchmarking. The main problem is when forming arrays, for example to generate (2,2) random normal array with double precision we should do:

import jax
jax.config.update('jax_enable_x64', True)
for backend in ['numpy', 'tensorflow', 'torch', 'jax', 'dask', 'mars', 'sparse']:
    if backend in ('tensorflow', 'torch'):
        A = ar.do("random.normal", size=(2,2), like=backend, dtype=ar.to_backend_dtype('float64', backend))
    else:
        A = ar.do("random.normal", size=(2,2), like=backend)

We can't just always supply the dtype argument, since numpy, dask and sparse throw an error when fed dtype. We could also generate whatever dtype array and then convert the result to double precision, but this doesn't really address the problem. This doesn't just hold for random.normal, but for essentially any kind of array-creating functions, like zeros or eye, although there supplying dtype does work. (for jax we still need to set jax_enable_x64.) I can also see from your gen_rand method in test_autoray.py that you encountered similar problems.

Suggested solutions

• Make a wrapper for numpy, dask, sparse (and cupy?) that ignores the dtype keyword, and then converts result to the correct dtype after the fact (if dtype is 'float32'). For jax we should maybe throw a warning if trying to generate double precision random numbers without setting 'jax_enable_x64' to True. In fact, for example for 'zeros', jax already throws a warning in this situation.
• Make a autoray.random.Generator object, like the numpy.random.Generator, but then backend aware. This may perform slightly better, and I think numpy is urging people to start using this over calling e.g. numpy.random.normal directly (although it doesn't seem to be catching on).

It might also be worthwhile to add translations for some more standard distributions like binomial or poisson, although I mostly use normal and uniform myself.

Hi, this is an interesting, yet difficult library! While it seems intriguing, I wonder thought about the interface: why is there always a 'do'? Why does the library not just wrap the API directly? Writing do("command") seems quite cumbersome compared to command.

Is that API a future plan?

In addition to the take translation I added in my previous PR, there is some more that might be good to add. At least, I am using these myself. I can make a PR.

• split. The syntax is different for numpy and tensorflow/torch. The former wants the number of splits or an array of locations of splits, whereas tensorflow/torch either want the number of splits or an array of split sizes. We can go from one format the other using np.diff
• diff. This is implemented in tensorflow as tf.experimental.numpy.diff, and not implemented at all for torch. This also means I don't know what the cleanest way is to implement split mentioned above. Maybe just using np.diff and then convert to array of right backend if necessary?
• linalg.norm, seems to work with tensorflow, but for torch we need to do _SUBMODULE_ALIASES["torch", "linalg.norm"] = "torch" I didn't check these things for any other libraries.

Maybe a bit of an overly ambitious idea, but have you ever thought about baking in support for JIT? Right now it seems that for TensorFlow everything works with eager execution, and I'm not sure you can compile the computation graphs resulting from a series of ar.do calls. PyTorch also support JIT to some extend with TorchScript Numpy doesn' t have JIT, but there is Numba Cupy has an interface with Numba that does seem to allow JIT. JAX has support for JIT

Another thing is gradients. Several of these libraries have automatic gradients, and having an autoray interface for doing computations with automatic gradients would be fantastic as well (although probably also ambitious).

If you think these things are doable at all, I wouldn't mind spending some time to try to figure out how this could work.

Less ambitiously, you did mention in #3 that something along the lines of

with set_backend(like):
    ...

would be pretty nice. I can try to do this. This probably comes down to checking for a global flag in ar.do after the line

if like is None: