我试图优化一个大量使用稀疏矩阵操作的代码，并尝试使用jax如下
import jax.scipy as jsp from jax.scipy import sparse
然而，当试图转换时，从另一个矩阵创建稀疏矩阵，如下所示
sHamil_ONE= sparse.csr_matrix(Hamil_multi_pol)
我收到这条信息
AttributeError: module 'jax.scipy.sparse' has no attribute 'csr_matrix'
那我该怎么办
m0rkklqb1#
JAX没有为scipy.sparse矩阵API提供 Package 器，但是jax.experimental.sparse为与jit、vmap、autodiff和其他JAX转换兼容的稀疏数组提供了一些实验性支持。用法如下所示：
scipy.sparse
jax.experimental.sparse
jit
vmap
from jax.experimental import sparse Hamil_ONE = sparse.BCOO.fromdense(Hamil_multi_pol)
你可以在这里阅读更多：https://jax.readthedocs.io/en/latest/jax.experimental.sparse.html
xzlaal3s2#
它说jax.scipy.sparse模块没有csr_matrix函数。相反，您可以使用jax.numpy：
jax.scipy.sparse
csr_matrix
jax.numpy
import jax.numpy as jnp # create a matrix with some zeros data = jnp.array([1, 2, 0, 0, 3, 4]) indixs = jnp.array([0, 1, 1, 2, 0, 1]) indptr = jnp.array([0, 2, 3, 6]) # create the csr_matrix (2x3) csr_matrix = jsp.sparse.csr_matrix((data, indixs, indptr), shape=(2, 3))
2条答案
按热度按时间m0rkklqb1#
JAX没有为
scipy.sparse矩阵API提供 Package 器，但是jax.experimental.sparse为与jit、vmap、autodiff和其他JAX转换兼容的稀疏数组提供了一些实验性支持。用法如下所示：
你可以在这里阅读更多：https://jax.readthedocs.io/en/latest/jax.experimental.sparse.html
xzlaal3s2#
它说
jax.scipy.sparse模块没有csr_matrix函数。相反，您可以使用jax.numpy：