Implement ufunc.outer

[eric-wieser]

Would be handy if np.add.outer worked

[mattjj]

I don't think our current primitive-wrapping strategy supports using any ufunc methods, but some of that functionality (like ufunc.outer) can be achieved with broadcasting, which we do support:

In [1]: a = b = np.arange(3)

In [2]: np.add.outer(a, b)
Out[2]:
array([[0, 1, 2],
       [1, 2, 3],
       [2, 3, 4]])

In [3]: a[:,None] + b
Out[3]:
array([[0, 1, 2],
       [1, 2, 3],
       [2, 3, 4]])

Is it possible you can use broadcasting for your use case? If not, can you provide an example of the kind of code you need to run with ufunc.outer?

It's possible we might be able to add ufunc method support easily by introducing a ufunc_primitive subclass to primitive, and implementing the ufunc methods using high-level things like broadcasting. I don't use ufunc methods, though, so I'm a bit in the dark on how much needs to be done.

[eric-wieser]

Is it possible you can use broadcasting for your use case?

Yes, absolutely. My request is that autograd do this for me. I'd be in favor of a ufunc_primitive class, although in 1.13, you don't need to override ufuncs at all, as __array_ufunc__ is a thing

As a first order approximation, this should suffice to implement outer:

def outer(ufunc, a, b, **kwargs):
	a = a.reshape(a.shape + (1,)*b.ndim)
	b = b.reshape((1,)*a.ndim + b.shape)
	return ufunc(a, b, **kwargs)

I don't use ufunc methods, though, so I'm a bit in the dark on how much needs to be done.

Right now, reduce_at is kinda broken in numpy anyway, and at doesn't make sense to support since you don't support A[i] = 0 anyway. So the only other functions to add are .reduce and accumulate, which are usually just sum/prod and cumsum/cumprod

[eric-wieser]

Elaborating on that ufunc comment, in 1.13 you can handle ufuncs as:

class ArrayNode(autograd.core.Node, np.lib.mixins.NDArrayOperatorsMixin):
    # ...
    def __array_ufunc__(self, ufunc, method, *args, **kwargs):
        p = get_the_autograd_primitive_for(ufunc)
        if method == '__call__:
            return p(*args, **kwargs)
        elif method == 'outer':
            do_the_above
        else:
            raise NotImplementedError('Request support for ufunc.{} on the issue tracker'.format(method))

    # no need to implement any operator overloads here

Leaving np.add completely unchanged

[j-towns]

I think this would be straighforward to implement in the way you describe above, with a ufunc_primitive class (or maybe call it binary_ufunc_primitive since the ufunc methods are only supported for ufuncs with two inputs). Then you could add an outer method to that class as you suggest. This would probably be the easiest to implement in terms of lines of code, although performance would not be optimal, since those calls to np.reshape etc would introduce a bit of overhead.

So just to be explicit, this should work (I'd probably put it in numpy_extra.py):

class binary_ufunc_primitive(autograd.core.primitive):
    def outer(self, a, b, **kwargs):
        a = a.reshape(a.shape + (1,)*b.ndim)
	b = b.reshape((1,)*a.ndim + b.shape)
	return self(a, b, **kwargs)

Then just have the binary ufuncs in numpy_grads.py and use that class.

[j-towns]

Otherwise you could do something like this (I find the subclassing syntax in Python pretty confusing so forgive me if this isn't correct):

class binary_ufunc_primitive(autograd.core.primitive):
    def __init__(self, fun):
        self.outer = autograd.core.primitive(self.fun.outer)
        super(binary_ufunc_primitive, self).__init__()

This would probably give you slightly better performance, but at some point you'd also need to define the vjp of outer by hand, which is more effort/lines of code.

[j-towns]

I reckon I could incorporate this (and the other binary ufunc methods) into #312.