Skip to content

Fix unit_diagonal and non-contiguous reads in MLX linalg - #2403

Open
jessegrabowski wants to merge 2 commits into
pymc-devs:mainfrom
jessegrabowski:mlx-solve-triangular-unit-diagonal
Open

Fix unit_diagonal and non-contiguous reads in MLX linalg#2403
jessegrabowski wants to merge 2 commits into
pymc-devs:mainfrom
jessegrabowski:mlx-solve-triangular-unit-diagonal

Conversation

@jessegrabowski

@jessegrabowski jessegrabowski commented Aug 27, 2026

Copy link
Copy Markdown
Member

MLX's solve_triangular has no unit_diagonal, and the dispatch dropped the flag rather than emulating it, so it solved against A's real diagonal. lu_solve hit this too, since it hands the packed LU factor straight in and relies on the flag to ignore U's diagonal.

Second, unrelated: mx.array keeps a non-C-contiguous numpy layout and MLX's elementwise kernels misread the buffer, so mlx_typify normalizes to C order. That made lu_solve's gradient wrong, and reverting it segfaults rather than failing on the numbers, so the read is going out of bounds.

Closes #2384

@jessegrabowski jessegrabowski added bug Something isn't working mlx linalg Linear algebra labels Aug 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working linalg Linear algebra mlx

Projects

None yet

Development

Successfully merging this pull request may close these issues.

MLX backend silently ignores unit_diagonal in SolveTriangular

1 participant