jax.experimental.sparse.bcsr_fromdense

内容

jax.experimental.sparse.bcsr_fromdense#

jax.experimental.sparse.bcsr_fromdense(mat, *, nse=None, n_batch=0, n_dense=0, index_dtype=<class 'jax.numpy.int32'>)[source]#

从稠密矩阵创建 BCSR 格式的稀疏矩阵。

参数：

mat (ArrayLike) – 要转换为 BCOO 的数组。
nse (int | None | None) – 每个批次中存储的元素数量
n_batch (int) – 批次维度数量（默认值：0）
n_dense (int) – 稠密维度数量（默认值：0）
index_dtype (DTypeLike) – 稀疏索引的数据类型（默认值：int32）

返回值：

矩阵的 BCSR 表示形式。

返回类型：

mat_bcsr