lineax

lineax copied to clipboard

New solver for block tridiagonal matrices

7

Hi, I needed a block tridiagonal solver for a project of mine, so I took a stab at adding one to lineax (very nice package, thank you!). The solve is...

aidancrilly

Implementation of LSMR for iterative least squares.

1

Tests seem to be passing locally. I added `LSMR` in the places it seemed to make sense (mainly places where I saw `SVD` which solves similar problems. Let me know...

f0uriest

Linear solvers for sparse matrices

4

This is just a question: are there any plans to support JAX-compatible linear solvers for sparse matrices? I am thinking of sparse linear systems of the type $Ax=b$ where $A$...

arpastrana

Iterative least squares solvers

1

Hi @patrick-kidger , I've really enjoyed playing around with Lineax! You have several iterative methods for square systems (CG, GMRES, BICGSTAB) but nothing for least squares problems. I have a...

f0uriest

Why is a non-materialised FunctionLinearOperator recompiled for every new input?

1

In [lineax.materialise() docs](https://docs.kidger.site/lineax/api/functions/#lineax.linearise) there is a following example: ```large_function = ... operator = lx.FunctionLinearOperator(large_function, ...) # Option 1 out1 = operator.mv(vector1) # Traces and compiles `large_function` out2 = operator.mv(vector2) #...

gawelk

Sometimes we detect that the output was `NaN`, and fail the linear solve. Under these circumstances we should check whether the input had any `NaN`s as well -- indicating that...

patrick-kidger

Batched JAX linear solves bugged for large batches

5

Hello, I opened a similar issue on the main JAX (https://github.com/google/jax/issues/19431) but I thought it may get more attention here. The batched JAX linear solves seem to be bugged for...

ma-gilles

Ambiguous output from successful solution using `lineax.RESULTS`

1

I originally encountered this issue in `optimistix` but realised it originates here. After a successful solution is found the `solution.result` value is `0` when it has succeeded, with a series...

LouisDesdoigts

How to use PyTreeLinearOperator

8

Hi lineax community, I came from Patrick's comment (https://github.com/google/jax/discussions/17203#discussioncomment-6814678). I believe `PyTreeLinearOperator` will do the job but am struggling to use it correctly. In the example below, I want to...

ToshiyukiBandai

Issue with `vmap` when using `lx.linear_solve` on `SparseMatrixOperator` with multi-column RHS

1

I'm encountering a ValueError when using the `vmap` functionality to map the `lx.linear_solve` operation across multiple columns of a right-hand side matrix (RHS) with a `SparseMatrixOperator`. I expect the `_solve_beta`...

Dong555