linear solve of tridiagonal system

[marcobonici]

When trying to invert a tridiagonal system and compiling that function with Reactant, I get an error.

Here the MWE

using Reactant
using LinearAlgebra

Reactant.set_default_backend("cpu")

# Minimal tridiagonal system
n = 5

dl = Float32[1, 1, 1, 1]
 d = Float32[4, 4, 4, 4, 4]
 du = Float32[1, 1, 1, 1]
 b = Float32[1, 2, 3, 4, 5]

# Trace all inputs
 dlR = Reactant.to_rarray(dl)
  dR = Reactant.to_rarray(d)
 duR = Reactant.to_rarray(du)
  bR = Reactant.to_rarray(b)

function solve_tridiagonal(dl, d, du, b)
    A = Tridiagonal(dl, d, du)
    return A \ b
end

println("=== Reactant Tridiagonal MWE ===")
try
    f = Reactant.@compile solve_tridiagonal(dlR, dR, duR, bR)
    y = f(dlR, dR, duR, bR)
    println("SUCCESS: ", y)
catch e
    println("FAIL: ", typeof(e))
    println(e)
end

here the output

=== Reactant Tridiagonal MWE ===
FAIL: TypeError
TypeError(:typeassert, "", Bool, TracedRNumber{Bool}(()))

[wsmoses]

@mofeing / @avik-pal since iirc you were looking at solvers recently?

[avik-pal]

jax has the ffi for cuda/rocm https://github.com/jax-ml/jax/blob/5a780a031efdf8463280d5b8e809ef718eb69e3d/jax/_src/lax/linalg.py#L2768-L2787. CPU will need a manual implementation. Maybe we make it a proper op itself in enzymejax

[marcobonici]

Hi @wsmoses @avik-pal is there any update on this issues?
If not, if there is some guidance, I am prone to try to implement stuff myself (but again, I will need guidance xD).