Mlx Array Convenience Initializer Macro#355
Mlx Array Convenience Initializer Macro#355robertmsale wants to merge 5 commits intoml-explore:mainfrom
Conversation
| ```swift | ||
| import MLX | ||
|
|
||
| let a = #mlx([[1, 2], [3, 4]]) |
There was a problem hiding this comment.
This is an interesting idea -- I guess using a macro is the only way to get compile-time checks. I am not sure on the naming: #mlx seems a bit vague. I don't suppose #MLXArray would work, would it? I suspect the name would clash in the macro definition. Maybe #mlxarray or #array?
There was a problem hiding this comment.
Good point! I was worried that #MLXArray wouldn't work, but I think because macros get expanded before the AST is built, I replaced all #mlx macros with #MLXArray and all the tests passed successfully. It's kinda cool because it now reads like the thing it's initializing! You just throw a # behind it and suddenly can do nested arrays with strong shape & type guarantees, so that worked out pretty well!
Proposed changes
Potential solution for #161.
This change introduces a new
#mlx(...)expression macro to improve array literal ergonomics forMLXArrayinitialization while keeping expansion behavior explicit and testable:#mlxas a public macro surface inMLX, so users can write:#mlx([[1, 2], [3, 4]])#mlx([[1, 2], [3, 4]], dtype: .int16).float32(avoids a trailing.asType(...)cast)..asType(...)for dynamic dtype expressions and dtypes that do not cleanly map to Swift literal typing (.float16,.bfloat16,.complex64,.float64,.bool).Tests and docs:
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes