mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Test for gtsv2 attr on cusparse.

Browse Source
This commit is contained in:
Qiao Zhang 2021-06-23 14:22:08 -07:00
parent 9493b315bc
commit 7cc277d634
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/lax/linalg.py
View File

@ -1361,7 +1361,7 @@ tridiagonal_solve_p.def_impl(
functools.partial(xla.apply_primitive, tridiagonal_solve_p))
tridiagonal_solve_p.def_abstract_eval(lambda dl, d, du, b, *, m, n, ldb, t: b)
# TODO(tomhennigan): Consider AD rules using lax.custom_linear_solve?
if cusparse is not None:
if cusparse is not None and hasattr(cusparse, "gtsv2"):
xla.backend_specific_translations['gpu'][tridiagonal_solve_p] = cusparse.gtsv2