Implement missing `ufunc.outer` functions for binary ops

[ricardoV94]

Please describe the purpose of filing this issue

https://numpy.org/doc/stable/reference/generated/numpy.ufunc.outer.html

[ricardoV94]

Can we obtain this with simple dimshuffles?

import numpy as np

x = np.arange(3)
y = np.arange(5)

np.add.outer(x, y)
# array([[0, 1, 2, 3, 4],
#        [1, 2, 3, 4, 5],
#        [2, 3, 4, 5, 6]])

x[:, None] + y[None, :] 
# array([[0, 1, 2, 3, 4],
#        [1, 2, 3, 4, 5],
#        [2, 3, 4, 5, 6]])

def outer(op, x, y):
  x_ = np.expand_dims(x, tuple(range(-y.ndim, 0)))
  y_ = np.expand_dims(y, tuple(range(x.ndim)))
  return op(x_, y_)

x = np.ones((3, 5))
y = np.ones((7, 1, 3))

np.testing.assert_array_equal(
  outer(np.add, x, y),
  np.add.outer(x, y)
)

[ricardoV94]

Would require adding an outer method to the binary Elemwise Ops that would implement the same PyTensor logic as the NumPy outer function above.

[ricardoV94]

The idea would be to add an outer method to Elemwise

pytensor/pytensor/tensor/elemwise.py

Line 305 in bfd7257

class Elemwise(OpenMPOp):

Roughly, it might look something like:

def outer(self, x, y):
  # Add and Mul scalar op nin is -1, because they allow arbitrary number of inputs
  if self.scalar_op.nin not in (-1, 2):
    raise NotImplementedError("outer is only available for binary operators")
  
  x_ = pt.expand_dims(x, tuple(range(-y.ndim, 0)))
  y_ = pt.expand_dims(y, tuple(range(x.ndim)))
  return self(x_, y_)

And then one should be able to call:

import pytensor.tensor as pt
z = pt.add.outer(x, y)

CC @Raj-Parekh24, let me know if this is sufficient to get you started.