Skip to content

Add MLX dispatch for CholeskySolve - #2377

Open
jessegrabowski wants to merge 2 commits into
pymc-devs:mainfrom
jessegrabowski:mlx-cho-solve
Open

Add MLX dispatch for CholeskySolve#2377
jessegrabowski wants to merge 2 commits into
pymc-devs:mainfrom
jessegrabowski:mlx-cho-solve

Conversation

@jessegrabowski

Copy link
Copy Markdown
Member

MLX has no cho_solve, so this runs two triangular solves against the Cholesky factor and its transpose.

Every op in the dispatch, including the transpose, is pinned to the CPU stream. mlx.linalg is CPU-only, and a transpose left on the default stream aborts on float64 input.

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant