Dynamic Shape Support

[curtischong]

I am currently using einx in a graph neural network, where the number of edges are dynamic between graphs.

import time, torch, einx

def run_einx():
    for E in range(1000, 1050):  # 50 novel shapes
        d, c = torch.randn(E, 8), torch.rand(E)
        einx.divide("e d, e -> e d", d, c)

def run_torch():
    for E in range(1000, 1050):
        d, c = torch.randn(E, 8), torch.rand(E)
        d / c.unsqueeze(-1)  # shape-agnostic: no per-shape cost

t0 = time.perf_counter(); run_einx(); print(f"einx, novel shapes: {time.perf_counter()-t0:.3f}s")
t0 = time.perf_counter(); run_einx(); print(f"einx, cached:       {time.perf_counter()-t0:.3f}s")
t0 = time.perf_counter(); run_torch(); print(f"raw torch:          {time.perf_counter()-t0:.3f}s")

This script produces this output:

einx, novel shapes: 0.481s
einx, cached:       0.004s
raw torch:          0.002s

I think the reason why einx is slower is because einx needs to generate a new expression every time (because the shapes of the tensors d and c are not consistent each time).

To see the code that's generated by einx, I ran

import torch, einx
dist = torch.rand(4096); centers = torch.linspace(0, 6, 32)
print(einx.subtract("e, r -> e r", dist, centers, graph=True))

Which generates this code:

import torch
def op(a, b):
    a = torch.reshape(a, (4096, 1))
    b = torch.reshape(b, (1, 32))
    c = torch.subtract(a, b)
    return c

I believe that these reshapes are generated for each unique call (because we need to hardcode the dimensions such as 4096). This means that graphs with different number of edges will not use the same einx code and will therefore be much slower as custom code needs to be generated each time to have the proper shapes.

I think we need some way to specify that some parameters are dynamic to trigger smarter compilation?

[fferflo]

einx currently assumes that all shapes are known at compile time (i.e. static), and does not support compiling operations with unknown dimension sizes. Supporting dynamic dimension would be great, but would require significant changes to einx internally, including:

• einx currently uses sympy to solve for all axis sizes. This would have to be changed such that the result is not a static int, but an equation for computing some axis size given the size of other axes.
• The tensor object in the computation graph currently assumes that its shape member is a tuple of constant ints. This would have to be changed to allow for a tuple of dynamic tracer objects instead (e.g. x.shape[0] // y.shape[1]).
• All operation implementations have to be checked to make sure that there aren't any that rely on statically known axis sizes.
• Some mechanism has to be added for deciding if a dimension should be assumed to be static or dynamic in compilation.

This might be added to einx in the future, but is currently not supported.

(As a sidenote: Some tensor frameworks such as Jax also only support JIT compilation for statically known axis sizes: https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#dynamic-shapes)