mlx-swift
mlx-swift copied to clipboard
ml-explore
Add arange
Proposed changes
Do you think it would make sense to add arange to mlx-swift? This is probably the most common piece of missing syntax that I've encountered when porting models from Python to Swift. There are enough existing examples that in practice it's easy enough to find a workaround, but having a direct equivalent would reduce friction. Is there a reason this was not included in mix-swift?
Checklist
- [x] I have read the CONTRIBUTING document
- [x] I have run
pre-commit run --all-filesto 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)
There are enough existing examples that in practice it's easy enough to find a workaround, but having a direct equivalent would reduce friction. Is there a reason this was not included in mix-swift?
It wasn't added because there is an idiomatic swift way of doing it (since Range and friends implement Collection):
let a = MLXArray(0 ..< 12)
There are range variants for open, closed and strides -- all of the arange calls can be expressed with these.
Do you think arange gives something compelling over these?
The main motivation would be discoverability, especially since models are now commonly ported with automated tools. This came to mind as I was cleaning up the mlx-audio repo, which contained some workarounds for arange in Swift.
Perhaps there's also a performance benefit of using the MLX C API directly, which uses lazy computation?
The syntax with step is also more concise:
// Current workarounds
let a = MLXArray(stride(from: 0, to: 12, by: 2))
let b = MLXArray(stride(from: 0.0, to: 5.0, by: 0.5))
let a = MLXArray.arange(0, 12, step: 2)
let b = MLXArray.arange(0.0, 5.0, step: 0.5)
That seems reasonable -- we can cross link the documentation between the two styles. Using mlx_arange may actually be a performance gain as it will probably initialize the array on the GPU.
