Fix shapeless matmul with dynamic batch dimensions

[varshneydevansh]

Proposed changes

• In shapeless dynamic tracing, avoid the existing flatten -> matmul -> unflatten
  optimization because unflatten stores the trace-time shape, and reuses it
  when the dynamic sequence length changes.

• Uses the existing batched/broadcasted matmul path during dynamic tracing so
  the output shape follows the runtime input shape.

• Keeps the existing flatten/unflatten path for non-dynamic tracing to preserve
  the current optimized behavior.

• Fixes [BUG] shapeless matmul isn't #2607

Tests:

• Added test_export_matmul_shapeless_mid_dim, based on the reproducer from
  [BUG] shapeless matmul isn't #2607.

• Added test_export_matmul_shapeless_batch_and_mid_dim, based on the coverage
  idea from fix(shapeless): infer dynamic dim in Unflatten::output_shapes (#2607) #3649
  but extended to check numerical correctness with nonzero inputs.

Checklist

Put an x in the boxes that apply.

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

[varshneydevansh]

Here are the test cases which I ran before the fix :

 Devanshs-MacBook-Pro    \mlx  ➜ (  shapeless_matmul)   2.896s   9:28 AM   
 ⚡devanshvarshney ❯❯ python3 -m pytest python/tests/test_export_import.py -k matmul_shapeless_mid_dim
======================================================================================= test session starts =======================================================================================
platform darwin -- Python 3.14.2, pytest-9.1.1, pluggy-1.6.0
rootdir: /Users/devanshvarshney/mlx
configfile: pyproject.toml
collected 19 items / 18 deselected / 1 selected                                                                                                                                                   

python/tests/test_export_import.py .                                                                                                                                                      [100%]

============================================================================================ FAILURES =============================================================================================
_______________________________________________________________ TestExportImport.test_export_matmul_shapeless_mid_dim (seq_len=248) _______________________________________________________________

self = <test_export_import.TestExportImport testMethod=test_export_matmul_shapeless_mid_dim>

    def test_export_matmul_shapeless_mid_dim(self):
        path = os.path.join(self.test_dir, "matmul_shapeless.mlxfn")
    
        E, H = 64, 17
        arr = mx.arange(E * H, dtype=mx.float32).reshape((E, H)) * (1.0 / (E * H))
    
        def fn(x):
            return mx.matmul(x, arr)
    
        sample = mx.zeros((1, 40, E), dtype=mx.float32)
        mx.export_function(path, fn, sample, shapeless=True)
        imported = mx.import_function(path)
    
        for seq_len in (40, 248, 623):
            with self.subTest(seq_len=seq_len):
                x = mx.arange(seq_len * E, dtype=mx.float32).reshape((1, seq_len, E))
                expected = fn(x)
                (y,) = imported(x)
>               self.assertEqual(y.shape, (1, seq_len, H))
E               AssertionError: Tuples differ: (1, 40, 17) != (1, 248, 17)
E               
E               First differing element 1:
E               40
E               248
E               
E               - (1, 40, 17)
E               ?      ^
E               
E               + (1, 248, 17)
E               ?     + ^

python/tests/test_export_import.py:725: AssertionError
_______________________________________________________________ TestExportImport.test_export_matmul_shapeless_mid_dim (seq_len=623) _______________________________________________________________

self = <test_export_import.TestExportImport testMethod=test_export_matmul_shapeless_mid_dim>

    def test_export_matmul_shapeless_mid_dim(self):
        path = os.path.join(self.test_dir, "matmul_shapeless.mlxfn")
    
        E, H = 64, 17
        arr = mx.arange(E * H, dtype=mx.float32).reshape((E, H)) * (1.0 / (E * H))
    
        def fn(x):
            return mx.matmul(x, arr)
    
        sample = mx.zeros((1, 40, E), dtype=mx.float32)
        mx.export_function(path, fn, sample, shapeless=True)
        imported = mx.import_function(path)
    
        for seq_len in (40, 248, 623):
            with self.subTest(seq_len=seq_len):
                x = mx.arange(seq_len * E, dtype=mx.float32).reshape((1, seq_len, E))
                expected = fn(x)
                (y,) = imported(x)
>               self.assertEqual(y.shape, (1, seq_len, H))
E               AssertionError: Tuples differ: (1, 40, 17) != (1, 623, 17)
E               
E               First differing element 1:
E               40
E               623
E               
E               - (1, 40, 17)
E               ?     ^^
E               
E               + (1, 623, 17)
E               ?     ^^^

python/tests/test_export_import.py:725: AssertionError
===================================================================================== short test summary info =====================================================================================
SUBFAILED(seq_len=248) python/tests/test_export_import.py::TestExportImport::test_export_matmul_shapeless_mid_dim - AssertionError: Tuples differ: (1, 40, 17) != (1, 248, 17)
SUBFAILED(seq_len=623) python/tests/test_export_import.py::TestExportImport::test_export_matmul_shapeless_mid_dim - AssertionError: Tuples differ: (1, 40, 17) != (1, 623, 17)
=========================================================================== 2 failed, 1 passed, 18 deselected in 1.11s ============================================================================

Here are test cases I ran after the fix :

 Devanshs-MacBook-Pro    \mlx  ➜ (  shapeless_matmul)   2m 7.641s   8:53 PM   
 ⚡devanshvarshney ❯❯ python3 -m pytest python/tests/test_export_import.py -k matmul_shapeless_mid_dim
======================================================================================= test session starts =======================================================================================
platform darwin -- Python 3.14.2, pytest-9.1.1, pluggy-1.6.0
rootdir: /Users/devanshvarshney/mlx
configfile: pyproject.toml
collected 19 items / 18 deselected / 1 selected                                                                                                                                                   

python/tests/test_export_import.py .                                                                                                                                                        [100%]

================================================================================ 1 passed, 18 deselected in 0.20s =================================================================================


 Devanshs-MacBook-Pro    \mlx  ➜ (  shapeless_matmul)   493ms   9:07 PM   
 ⚡devanshvarshney ❯❯ python3 -m pytest python/tests/test_export_import.py

======================================================================================= test session starts =======================================================================================
platform darwin -- Python 3.14.2, pytest-9.1.1, pluggy-1.6.0
rootdir: /Users/devanshvarshney/mlx
configfile: pyproject.toml
collected 19 items                                                                                                                                                                                

python/tests/test_export_import.py ...................                                                                                                                                      [100%]

======================================================================================= 19 passed in 1.31s ========================================================================================


 Devanshs-MacBook-Pro    \mlx  ➜ (  shapeless_matmul)   283ms   9:26 PM   
 ⚡devanshvarshney ❯❯ python3 -m pytest python/tests/test_export_import.py -k matmul_shapeless_batch_and_mid_dim
======================================================================================= test session starts =======================================================================================
platform darwin -- Python 3.14.2, pytest-9.1.1, pluggy-1.6.0
rootdir: /Users/devanshvarshney/mlx
configfile: pyproject.toml
collected 20 items / 19 deselected / 1 selected                                                                                                                                                   

python/tests/test_export_import.py .                                                                                                                                                        [100%]

================================================================================ 1 passed, 19 deselected in 0.14s =================================================================================

[zcbenz]

I think your fix does the work correctly and we were probably overthinking the problem. Need confirmation from @angeloskath.

[varshneydevansh]

Yes, awni's comment already had the key hint. While I was studying
how the flatten/unflatten and broadcasting works it clicked me and
when I looked at the code voila the solution was right there :)