Open davidweichiang opened 2 years ago
Yes, this would be a good feature to add, and I don't think it would be too tough to implement.
Maybe just call torch.expand on the arguments beforehand?
Or is it possible that it just works already, but the dimension checks are too strict?
torch.einsum now supports broadcasting: https://github.com/pytorch/pytorch/issues/48420