Support shapeless compilation of scan operations

[keeeeenw]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure:
• Codex with GPT6-Astra to learn about the overall project and suggest changes especially for the test.

Hi,

I’m an ML engineer and researcher, a passionate Apple user. I’d like to contribute to MLX and develop a deeper understanding of its internals. I’m particularly excited about its potential for local inference and building LLMs end to end, from pretraining to post-training.

This pull request fixes #4460.

This adds DEFINE_INPUT_OUTPUT_SHAPE() to Scan, allowing it to infer the output shape from its input.

Please note that I also noticed https://github.com/ml-explore/mlx/pull/4468/changes has a similar fix but was rejected due to use of AI. Please advise me on the best practices for this type of fixes if this fix does not meet your standards! Also, please let me know if there is any better starter bug/task for me. I am happy to help and learn more during the process!

Validation

Tested on my Apple silicon (M1 Max with 32G of unified memory) with macOS 26.6.2 and Python 3.12.10:

• Before the fix, the newly added test_shapeless_compile_scan raised the missing output-shape inference error when tested with either CPU or GPU. After the fix, all pass on both devices.
• All CPU compilation tests and CPU scan tests passed.
• All repository pre-commit checks passed.