Tensor.einsum does not correctly handle unspecified dimensions

[CamelCaseCam]

I've been converting a model from PyTorch to tinygrad, and ran into an issue with einsums (defined in Tensor.py).

The einsum "...id, ...jd -> ...ij" executes fine in PyTorch but throws an error (below) when executed in tinygrad on the exact same input data. I've included a minimal example below the error for debugging.

import tinygrad
import torch
import numpy as np

# Define the input tensor
inp1_np = np.random.randn(16, 29, 256).astype(np.float32)
inp2_np = np.random.randn(16, 29, 256).astype(np.float32)

# Convert the input tensors to torch tensors
inp1 = torch.tensor(inp1_np)
inp2 = torch.tensor(inp2_np)

out = torch.einsum("...id, ...jd -> ...ij", inp1, inp2)
# Evaluate the output tensor
out_np = out.numpy()

# Convert the input tensors to tinygrad tensors
inp1_tg = tinygrad.Tensor(inp1_np)
inp2_tg = tinygrad.Tensor(inp2_np)

# Perform the einsum operation using tinygrad - will throw error!
out_tg = tinygrad.Tensor.einsum("...id, ...jd -> ...ij", inp1_tg, inp2_tg)

# Evaluate the output tensor
out_tg_np = out_tg.numpy()

I looked at the locals when the error appeared, and it looks like the "..." part of the einsum is being interpreted as an index for each of the dots. I might be wrong, but I think this is supposed to represent an arbitrary number of indices, not specific indices.

Let me know if you need any more information!