Add xtensor broadcast

[AllenDowney]

This replaces #1486. This one is based on a rebased labeled_tensor branch

---

📚 Documentation preview 📚: https://pytensor--1489.org.readthedocs.build/en/1489/

[AllenDowney]

@ricardoV94 Here's my attempt to rebase on the changes you just force pushed. Looks like mypy is unhappy -- is that something you expected?

Other than that, I think this is ready for review.

[ricardoV94]

@ricardoV94 Here's my attempt to rebase on the changes you just force pushed. Looks like mypy is unhappy -- is that something you expected?

Other than that, I think this is ready for review.

Yeah I didn't make mypy pass yet

[ricardoV94]

This won't work when the output shape is not statically known. The target shape has to be computer symbolically from the symbolic input shapes.

You can test by having an xtensor with shape=(None) for a dim that only that tensor has

[AllenDowney]

@ricardoV94 I think I have symbolic dimensions working. My solution is more complicated than I think any of us would like, but I don't see a simpler solution. Maybe you will.

Should we continue work on this PR, for now, and I will rebase later?

[ricardoV94]

Here is an idea:

def lower_broadcast(fgraph, node):
  excluded_dims = node.op.exclude
  broadcast_dims = tuple(dim for dim in node.outputs[0].type.dims if dim not in excluded_dims)
  all_dims = broadcast_dims + excluded_dims
  
  # align inputs with all_dims like we do in other rewrites 
  # probably time to refactor this kind of logic into a helper
  inp_tensors = []
  for inp, out in zip(node.inputs, node.outputs, strict=True)
    inp_dims = inp.type.dims
    order = tuple(inp_dims.index(dim) if dim in inp_dims else "x" for dim in all_dims)
    inp_tensors.append(inp.values.dimshuffle(order))
  
  if not excluded_dims:
    out_tensors = pt.broadcast_arrays(*inp_tensors)
  else:
    all_shape = tuple(pt.broadcast_shape(*inp_tensors))
    assert len(all_shape) == len(all_dims)
    for inp_tensor, out in zip(inp_tensors, node.outputs):
      out_dims = out.type.dims
      out_shape = tuple(length for length, dim in zip(all_shape, all_dims) if dim in out_dims)
      out_tensors.append(pt.broadcast_to(inp_tensor, out_shape)

  new_outs = [as_xtensor(out_tensor, dims=out.type.dims) for out_tensor, out in zip(out_tensors, node.outputs)]
  return new_outs

Btw the base branch is merged. You can rebase/ start from it. Note that you don't need to open a new PR. You can force-push your changes after cleaning up the branch to your current remote

[AllenDowney]

@ricardoV94 I've added broadcast_like. Once we're happy with broadcast and broadcast_like, I will factor out some common code.

[AllenDowney]

Your version of lower_broadcast works when exclude is empty, but it fails on tests that have excluded dims.

I'll work on debugging it, but at the moment it's not clear to me whether these is a small error in your implementation or an actual problem with the logic.

[ricardoV94]

I suspect some wrong assumption on the excluded dims alignment but the general idea should work

[AllenDowney]

I think the incorrect assumption is that all outputs have the same shape. When exclude is not empty, they don't, in general.

[ricardoV94]

Actually there's a logical flaw. Two inputs could have an excluded dim with the same name but different length, in which case they shouldn't be aligned for the broadcast shape.

We should add that as a test.

Still the logic for each output should be something like broadcast_to(tensor, common_broadcast_shape + original_excluded_shape)

[ricardoV94]

I think the incorrect assumption is that all outputs have the same shape. When exclude is not empty, they don't, in general.

I didn't assume that, the dimshuffle was supposed to take care of that so that things were put in different axis for broadcasting. Still as I just wrote there was a wrong assumption that you could align shared excluded dims. They don't even come out in a uniform order do they?

[ricardoV94]

I don't think this logical flaw is why the tests are failing though. We should test that case as well

[ricardoV94]

@AllenDowney does this work? If not, what case fails?

from pytensor.xtensor.type import xtensor
from pytensor.xtensor.math import second

def broadcast(array, *arrays, exclude=()):
    if isinstance(exclude, str):
        exclude = (exclude,)
        
    def sum_excluded_dims(array):
        if not exclude:
            return array
        dims = array.dims
        array_exclude = tuple(e for e in exclude if e in dims)
        if not array_exclude:
            return array
        return array.sum(array_exclude)
    
    def align_excluded_dims(array):
        if not exclude:
            return array
        dims = array.dims
        array_exclude = tuple(e for e in exclude if e in dims)
        if not array_exclude:
            return array
        return array.transpose(..., *array_exclude)
        
    if not arrays:
        return array
    
    # Find broadcast shape by doing nested second after excluding via `sum`
    # The sum operation will be removed in rewrites, since only the shape matters
    broadcast_array = sum_excluded_dims(array)
    for other_array in arrays:
        # second is equivalent no `np.broadcast_arrays(x, y)[1]`
        broadcast_array = second(broadcast_array, sum_excluded_dims(other_array))

    # Broadcast each original array with the broadcast_array
    # We further align the excluded dims according to the order given by the user, like xarray does
    return tuple(second(broadcast_array, align_excluded_dims(arr)) for arr in (array, *arrays))

    
x = xtensor(dims=("a", "b", "c"))
y = xtensor(dims=("a", "d"))
z = xtensor(dims=("a", "f", "b"))
for out in broadcast(x, y, z, exclude=("b", "f")):
    print(out.dims)
# ('a', 'c', 'd', 'b')
# ('a', 'c', 'd')
# ('a', 'c', 'd', 'b', 'f')