correctly infer backend for functions like concatenate

[Rik Voorhaar]

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.

[Johnnie Gray]

Yes this is a good point that I've run into a few times. Regarding your approaches:

1. I think this would be quite nice. Since such a dispatcher would just need the name of the function I'm not sure that that it would even need to be exposed to any users of register_function.

2. While this might work in some cases, the problem here might be the caching - _infer_class_backend_cached is cached just on obj.__class__, so builtins can't be dispatched separately depending on the inner type of iterable, or vs not iterable at all. One could add logic in infer_backend but for performance reasons (this is called potentially every operation) this is not desirable. Moreover, for einsum the problem is the first arg is (usually) a string and not to be dispatched on anyway so one needs extra logic anyway.

---

A dispatcher might look like:

def einsum_dispatcher(*args, **kwargs):
    if isinstance(args[0], str):
        return args[1]
    return args[0]

register_dispatch('einsum', einsum_dispatcher)

and then in do something like:

    if like is None:
        dispatch_arg = _DISPATCHERS.get(fn, default_dispatcher)(*args, **kwargs)
        backend = infer_backend(dispatch_arg)

with

def default_dispatcher(*args, **kwargs):
    return args[0]

hopefully just having the one extra dict lookup woudn't have too much overhead, possibly a defaultdict(lambda: default_dispatcher) would help here too.

---

On the manually specifying side I've been wondering if something like

with set_backend(like):
    ...

to force a backend for a block would be sufficiently convenient to add, but the above solution I think is cleaner for this particular problem.

Depending on your thoughts I'm very happy to implement the custom dispatcher approach or would happily accept a PR along those lines, let me know.

[Rik Voorhaar]

A good point, my second suggestion will indeed not work.

I'll try and see if I can get the custom dispatcher approach to work. It doesn't seem like it would be too difficult.
Something along the lines of with set_backend(like) would also be really nice, and I think one can implement both.

I really like this project, so I would love to contribute. I use this extensively in a project of mine, and doing a PR is a nice excuse to dive a bit deeper into the workings of your code. I also have a bunch of extra translations I needed for my project, but I will come back to that once the project is nearer to being finished.

[Johnnie Gray]

Resolved by #4. Glad autoray has been useful and thanks again for the PR.