Fix einsum not broadcasting batch dimensions in batched tensordot

[Adityaj0]

Fixes #4124.

Proposed changes

batch_tensordot broadcast the contracting dimensions of its two operands to a common size but left the batch dimensions alone, then built the output shape from the first operand only:

for (auto ax : a_batch) {
  out_shape.push_back(a.shape(ax));
}

When the first operand had a size 1 batch dimension, matmul broadcast it against the second operand while out_shape kept the 1, so the final reshape got the wrong element count:

mx.einsum("...ij,...jk->...ik", mx.zeros((1, 3, 4)), mx.zeros((2, 4, 5)))
# ValueError: [reshape] Cannot reshape array of size 30 into shape (1,3,5).

The reverse operand order worked, because the output shape was then taken from the larger operand. That asymmetry is why it went unnoticed. It is not specific to ellipses either, "bij,bjk->bik" with the same shapes fails identically.

The fix broadcasts the batch dimensions the same way the contracting dimensions already are, in the same block. a_batch and b_batch are parallel by construction (b_batch is built by looking up each a_batch label in b's subscript), so they can be zipped exactly like a_contract / b_contract.

operands	NumPy	before	after
(1,3,4), (2,4,5)	(2,3,5)	error	(2,3,5)
(2,3,4), (1,4,5)	(2,3,5)	(2,3,5)	(2,3,5)
(1,4), (5,4) -> ...	(5,)	error	(5,)
(5,4), (1,4) -> ...	(5,)	(5,)	(5,)

Tests

Added test_ellipses_broadcast to python/tests/test_einsum.py covering both operand orders, multiple leading batch dimensions, a reduced (->...) output, and the explicit-label spelling, each checked against NumPy for both shape and values.

I also removed an unconditional return in test_broadcasting that made the four lines after it dead code. To be clear about scope: those assertions pass on main as well, so that is an unrelated cleanup rather than something this fix enables. I checked before touching it. It was simply hiding coverage right next to the bug.

• I have read the CONTRIBUTING document
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (not needed, no API change)
• I have run pre-commit run --all-files to format my code and installed pre-commit prior to committing changes

Verified with a CPU-only build (-DMLX_BUILD_METAL=OFF): test_einsum, test_ops, test_autograd, test_vmap, test_linalg, test_blas, test_nn, test_array and test_reduce pass (464 tests).

[JasonHonKL]

Will this cause some performance issue ? @Adityaj0

[Adityaj0]

Hi @JasonHonKL no, because the fix only adds lazy batch dimension broadcasting in the existing batched tensordot path, so it should have negligible overhead beyond a small amount of shape bookkeeping.