ENH: Symbolic Vectorization

[ferrine]

fg = FunctionGraph([inputs], [outputs])
# follow the JAX API https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html
vfg = pytensor.graph.vectorize(fg, in_axis=[0], out_axis=[0], axis_name={0: "batch"})

# additionally, some collective Ops seem to be useful
...
# no OP if there until pytensor.graph.vectorize is called
sum_batch = pytensor.tensor.collective.Sum(tensor, axis="batch")
mean_batch = pytensor.tensor.collective.Mean(tensor, axis="batch")
...

Context for the issue:

Graph rewriting to vectorize operations in a symbolic way is a huge step to improve pymc/pytensor user experience.

Example use cases:

• Remove scan from VI
  https://github.com/pymc-devs/pymc/blob/main/pymc/variational/opvi.py#L1355
• Vectrorize NUTS over chains, to be done and to be implemented

[ricardoV94]

Don't forget to see the work that was already started in aesara-devs/aesara#1215

CC @Sayam753 and @purna135

[ferrine]

Seems like that work is quite different from the proposal

[ricardoV94]

Seems like that work is quite different from the proposal

It sounds you're talking about a high-level interface whereas the Blockwise tackles one possible implementation at the Op level?

That PR implements Symbolic vectorization for individual Ops. What you're talking about would be something like a fusion of multiple vectorized Ops, the same way the Composites fuses multiple Elemwise?

In the end, vectorization is just an abstraction for looping with repeated operations over some tensors. If Scan was fast that wouldn't be a terrible first way to achieve vectorization.

---

Regarding NUTS, you can't really vectorize something like that unless you implement the algorithm in the same language.

[ferrine]

I thought about creating a new vectorized fgraph by applying vectorized translations.

an Op would have a dispatch rule

@_vectorize.register(Op):
def vectorize_node(op: Op, inputs: List[Variable], dims: List[Optional[int]]) -> List[List[Variable], List[Optional[int]]]:
    ...

This can potentially be multiple additional operations to achieve broadcasting patterns or completely different ops implementing the same thing

[ferrine]

Here is an draft implementation

import pytensor.tensor as pt
import pytensor.graph.fg
import pytensor.utils
import functools
from typing import *

@functools.singledispatch
def vectrorize_symbolic(op, inputs: List[pt.Variable], dims: List[Optional[int]]) -> (List[pt.Variable], List[Optional[int]]):
    if any(d is not None for d in dims):
        raise RuntimeError(f"vectorization was not implemented for {op}")
    return no_vectorize(op, inputs)

def no_vectorize(op, inputs: List[pt.Variable]) -> (List[pt.Variable], List[Optional[int]]):
    outs = pytensor.utils.flatten(op(*inputs))
    return outs, [None] * len(outs)

def batch_axis_at(tens, batch_axis, at=0):
    pattern = list(range(tens.ndim))
    if batch_axis is None:
        pattern.insert(0, "x")
    elif batch_axis == 0:
        return tens
    else:
        pattern.insert(at, pattern.pop(batch_axis))
    return tens.dimshuffle(*pattern)

@vectrorize_symbolic.register(pt.basic.Elemwise)
def vectorize_elementwise(op, inputs: List[pt.Variable], dims: List[Optional[int]]) -> (List[pt.Variable], List[Optional[int]]):
    # use broadcasting
    vec_inputs = [batch_axis_at(v, d) for v, d in zip(inputs, dims)]
    vec_outputs = pytensor.utils.flatten(op(*vec_inputs))
    return vec_outputs, [0]*len(vec_outputs)

@vectrorize_symbolic.register(pt.basic.DimShuffle)
def vectorize_dimshuffle(op, inputs: List[pt.Variable], dims: List[Optional[int]]) -> (List[pt.Variable], List[Optional[int]]):
    # move batch axis to first, rest axis are incremented
    def move(i, b):
        if b is None:
            return i
        elif i == "x":
            return "x"
        # [0, 1,  +1> 2, 3, 4]
        # [0, 1, "b", 2, 3, 4]
        elif i<b:
            return i
        elif i >=b:
            return i + 1
    batch_dim = dims[0]
    order = [0] + [move(i, batch_dim) for i in op.new_order]
    broadcastable = list(op.input_broadcastable)
    broadcastable.insert(batch_dim, False)
    new_op = pt.basic.DimShuffle(broadcastable, order)
    vec_outputs = pytensor.utils.flatten(new_op(*inputs))
    return vec_outputs, [0]

You can recreate vectorized fgraph in a toposort loop

x = pt.vector()
y = pt.vector()

z = (x + y) ** 2

inputs = [x, y]
outputs = [z.T]
fg = pytensor.graph.fg.FunctionGraph(inputs, outputs)
vec_mapping = dict(zip([pt.matrix(), pt.vector()], [0, None]))
inputs_map = dict(zip(fg.inputs, vec_mapping))


for apply in fg.toposort():
    orig_inputs = apply.inputs
    # try get a vectorized input for this operation
    # otherwise no batching was performed for the input
    vec_inputs = [inputs_map.get(o, o) for o in orig_inputs]
    vec_patterns = [vec_mapping.get(v, None) for v in vec_inputs]
    if all(d is None for d in vec_patterns):
        vec_outputs, dims = no_vectorize(apply.op, vec_inputs)
    else:
        vec_outputs, dims = vectrorize_symbolic(apply.op, vec_inputs, vec_patterns)
    vec_mapping.update(zip(vec_outputs, dims))
    inputs_map.update(zip(apply.outputs, vec_outputs))

vfg = pytensor.graph.fg.FunctionGraph([inputs_map[o] for o in fg.inputs], [inputs_map[o] for o in fg.outputs])

[ricardoV94]

Precisely, now if you pair this with Blockwise you expand coverage of Ops that can be vectorized to nearly 100%

For now you would be limited to Dimshuffles and Elemwises.

That's what I meant with keep that in mind.

[ferrine]

I see

[ferrine]

With this approach, we can be generic about how we rewrite, for blockwise ops we can just move dims around and this core loop is fixed. The two PRs compliment each other in a sense

[ricardoV94]

Closed via #306