not working with jax (and cupy)

[pedrozudo]

Hi,

I was playing around a little bit your great library. I found that it does not work with jax, and also not cupy, but in the readme it says it would. Just wondering whether you have dropped support for those libraries or whether somethgin is going amiss. Minimal example to reproduce the error is basically the example in the readme but but with jax instead of torch.

import jax.numpy as jnp

# Identity matrix
n_rows = 20
rows = jnp.arange(n_rows)
cols = jnp.arange(n_rows)
data = jnp.ones(n_rows)

solver = CholeskySolverF(n_rows, rows, cols, data, MatrixType.COO)

b = jnp.ones(n_rows)
x = jnp.zeros_like(b)

solver.solve(b, x)
# b = [1, ..., 1]

and here is the error:

solver = CholeskySolverF(n_rows, rows, cols, data, MatrixType.COO)
TypeError: __init__(): incompatible function arguments. The following argument types are supported: 1. __init__(self, n_rows: int, ii: ndarray[dtype=int32, shape=(*), order='C'], jj: ndarray[dtype=int32, shape=(*), order='C'], x: ndarray[dtype=float64, shape=(*), order='C'], type: cholespy.MatrixType) -> None

Invoked with types: cholespy.CholeskySolverF, int, jaxlib.xla_extension.ArrayImpl, jaxlib.xla_extension.ArrayImpl, jaxlib.xla_extension.ArrayImpl, cholespy.MatrixType

A similar error also happens with cupy.

[wjakob]

Can you double check that your input data has the expected dimensionality, contiguity, and data type (please see the type annotations in the error message).

[pedrozudo]

I think everything seems to be on order. Just to double check, I initialized everything with jax, mapped it to torch using dlpack and ran the solver with torch tensors. This seems to be working fine.

import jax
import jax.numpy as jnp
import torch

from cholespy import CholeskySolverF, MatrixType


def jax2torch(j):
    t = torch.utils.dlpack.from_dlpack(j)
    return t


def torch2jax(t):
    j = jax.dlpack.from_dlpack(t)
    return j


def main(n_rows, rows, cols, data, m_type, b, x, backend):

    if backend == "torch":
        rows = jax2torch(rows)
        cols = jax2torch(cols)
        data = jax2torch(data)

        b = jax2torch(b)
        x = jax2torch(x)

    try:
        solver = CholeskySolverF(
            n_rows,
            rows,
            cols,
            data,
            m_type,
        )
    except:
        print(f"creating solver with {backend} did not work")
        return

    try:
        solver.solve(b, x)
        x = torch2jax(x)

    except:
        print(f"solving with {backend} did not work")
        return

    return x


if __name__ == "__main__":
    n_rows = 20
    rows = jnp.arange(n_rows)
    cols = jnp.arange(n_rows)
    data = jnp.ones(n_rows)
    m_type = MatrixType.COO

    b = jnp.ones(n_rows)
    
    x = jnp.zeros_like(b)

    assert "gpu" in x.device.platform

    x_j = main(n_rows, rows, cols, data, m_type, b, x, "jax")
    print(f"Jax solution: {x_j}")

    x_t = main(n_rows, rows, cols, data, m_type, b, x, "torch")
    print(f"Torch solution: {x_t}")