you like tinygrad? you know einsum notation? than you could use tinsor

This tensor frontend will take care of all the unsqueese / permutes / expand you need

from tinsor import Shape, Tensor

S, T, U, V = Shape(S=5, T=3, U=4, V=6)

x = Tensor.ones(S, T) # (S=5, T=3)
y = Tensor.rand(T, U) # (T=3, U=4)

p = x * y # (S=5, T=3, U=4)

print(p.numpy().shape) # (5, 3, 4)

install

pip install -e .

formerly known as tinsorchflow.