Skip to content

Mlx Array Convenience Initializer Macro#355

Open
robertmsale wants to merge 5 commits intoml-explore:mainfrom
robertmsale:mlx-array-convenience-initializer
Open

Mlx Array Convenience Initializer Macro#355
robertmsale wants to merge 5 commits intoml-explore:mainfrom
robertmsale:mlx-array-convenience-initializer

Conversation

@robertmsale
Copy link
Contributor

Proposed changes

Potential solution for #161.

This change introduces a new #mlx(...) expression macro to improve array literal ergonomics for MLXArray initialization while keeping expansion behavior explicit and testable:

  • Adds #mlx as a public macro surface in MLX, so users can write:
    • #mlx([[1, 2], [3, 4]])
    • #mlx([[1, 2], [3, 4]], dtype: .int16)
  • Validates nested literals are rectangular and numeric, and emits diagnostics for ragged/invalid input.
  • Supports deep nested literals and shape inference during expansion.
  • Implements dtype-aware expansion optimization:
    • Emits typed literal initialization directly for known integer dtypes and .float32 (avoids a trailing .asType(...) cast).
    • Falls back to .asType(...) for dynamic dtype expressions and dtypes that do not cleanly map to Swift literal typing (.float16, .bfloat16, .complex64, .float64, .bool).
  • Defines mixed literal promotion semantics:
    • If any literal element is floating-point, expansion uses floating-point construction for the full array.
    • Provides warning when float is used and explicit integer dtype is specified

Tests and docs:

  • Adds macro unit tests covering:
    • integer, float, mixed, deep-nested literals
    • dtype-specific expansion behavior
    • fallback behavior for dynamic/non-native dtypes
    • ragged-literal diagnostics
    • compiler warning when mixing float literals with integer dtypes
  • Updates DocC initialization docs with a macro section, examples, and mixed-literal promotion notes.

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)

```swift
import MLX

let a = #mlx([[1, 2], [3, 4]])
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants