MLX Backend

[lkarthee]

Issue for tracking and coordinating mlx backend work:

mlx.math

• [ ] fft
• [ ] fft2
• [ ] rfft
• [ ] irfft
• [ ] stft
• [ ] istft
• [x] logsumexp #19578
• [ ] qr
• [ ] segment_sum #19652
• [ ] segment_max #19652
• [x] erfinv #19628

mlx.numpy

• [ ] einsum
• [ ] bincount
• [ ] nonzero
• [ ] cross
• [ ] vdot
• [ ] nan_to_num
• [ ] copy
• [ ] roll
• [x] median #19568 #19574
• [x] meshgrid #19574
• [x] conjugate
• [x] arctan2 #19759
• [ ] quantile
• [ ] imag
• [ ] real
• [ ] select
• [x] argpartition https://github.com/keras-team/keras/pull/19680
• [ ] slogdet
• [ ] select
• [ ] vectorize
• [ ] correlate
• [x] diag #19714
• [x] diagonal #19714

mlx.image

• [x] rgb_to_grayscale #19609
• [x] resize - #19699

mlx.nn

• [ ] max_pool
• [ ] avg_pool
• [ ] conv
• [ ] depthwise_conv
• [ ] separable_conv
• [ ] conv_transpose
• [ ] ctc_loss

mlx.rnn

• [ ] rnn
• [ ] lstm
• [ ] gru

mlx.linalg

• [x] cholesky
• [ ] det
• [ ] eig
• [ ] eigh
• [x] inv
• [ ] lu_factor
• [x] norm #19698
• [x] qr
• [ ] solve
• [ ] solve_triangular
• [x] svd

mlx.core

• [x] np.ndarray of i64 is being cast to i32 in mlx during conversion if dtype is not passed
• [x] https://github.com/ml-explore/mlx/issues/1076
• [ ] https://github.com/ml-explore/mlx/issues/1075
• [x] https://github.com/ml-explore/mlx/issues/1066
• [x] https://github.com/ml-explore/mlx/issues/1065