[BUG] RuntimeError: einsum(): the number of subscripts in the equation (4) does not match the number of dimensions (2) for operand 0 and no ellipsis was given

[kyegomez]

Describe the bug
A clear and concise description of what the bug is and what the main root cause error is. Test very thoroughly before submitting.

Traceback (most recent call last):
  File "/Users/defalt/Desktop/Athena/research/BitNet/bitnet/main.py", line 189, in <module>
    y = layer(x)
        ^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/defalt/Desktop/Athena/research/BitNet/bitnet/main.py", line 177, in forward
    x = attn(q, k, v, mask=mask) + x
        ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/zeta/nn/attention/attend.py", line 287, in forward
    dots = einsum(f"b h i d, {kv_einsum_eq} -> b h i j", q, k) * scale
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/torch/functional.py", line 378, in einsum
    return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: einsum(): the number of subscripts in the equation (4) does not match the number of dimensions (2) for operand 0 and no ellipsis was given

To Reproduce
Steps to reproduce the behavior:

1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error

Expected behavior
A clear and concise description of what you expected to happen.

Screenshots
If applicable, add screenshots to help explain your problem.

Additional context
Add any other context about the problem here.