jax.experimental.sparse.BCOO

jax.experimental.sparse.BCOO#

class jax.experimental.sparse.BCOO(args, *, shape, indices_sorted=False, unique_indices=False)[source]#

JAX 中实现的实验性批量 COO 矩阵

参数:

  • (data – 批量 COO 格式的数据和索引。

  • indices) – 批量 COO 格式的数据和索引。

  • shape (Shape) – 稀疏数组的形状。

  • args (tuple[Array, Array])

  • indices_sorted (bool)

  • unique_indices (bool)

data#

形状为 [*batch_dims, nse, *dense_dims] 的 ndarray，包含稀疏矩阵中显式存储的数据。

类型:

Array

indices#

形状为 [*batch_dims, nse, n_sparse] 的 ndarray，包含显式存储数据的索引。重复的条目将被求和。

类型:

Array

示例

从稠密数组创建稀疏数组

>>> M = jnp.array([[0., 2., 0.], [1., 0., 4.]])
>>> M_sp = BCOO.fromdense(M)
>>> M_sp
BCOO(float32[2, 3], nse=3)

检查内部表示

>>> M_sp.data
Array([2., 1., 4.], dtype=float32)
>>> M_sp.indices
Array([[0, 1],
       [1, 0],
       [1, 2]], dtype=int32)

从稀疏数组创建稠密数组

>>> M_sp.todense()
Array([[0., 2., 0.],
       [1., 0., 4.]], dtype=float32)

从 COO 数据和索引创建稀疏数组

>>> data = jnp.array([1., 3., 5.])
>>> indices = jnp.array([[0, 0],
...                      [1, 1],
...                      [2, 2]])
>>> mat = BCOO((data, indices), shape=(3, 3))
>>> mat
BCOO(float32[3, 3], nse=3)
>>> mat.todense()
Array([[1., 0., 0.],
       [0., 3., 0.],
       [0., 0., 5.]], dtype=float32)
__init__(args, *, shape, indices_sorted=False, unique_indices=False)[source]#
参数:

方法

__init__(args, *, shape[, indices_sorted, ...])

astype(*args, **kwargs)

复制数组并转换为指定的数据类型。

block_until_ready()

from_scipy_sparse(mat, *[, index_dtype, ...])

从 scipy.sparse 数组创建 BCOO 数组。

fromdense(mat, *[, nse, index_dtype, ...])

从（稠密）Array 创建 BCOO 数组。

reshape(*args, **kwargs)

返回一个包含相同数据但具有新形状的数组。

sort_indices()

返回索引已排序的矩阵副本。

sum(*args, **kwargs)

沿轴对数组求和。

sum_duplicates([nse, remove_zeros])

返回已对重复索引求和的数组副本。

todense()

创建数组的稠密版本。

transpose([axes])

创建一个包含转置的新数组。

tree_flatten()

tree_unflatten(aux_data, children)

update_layout(*[, n_batch, n_dense, ...])

更新 BCOO 矩阵的存储布局（即 n_batch 和 n_dense）。

属性

T

dtype

n_batch

n_dense

n_sparse

ndim

nse

size

data

indices

shape

indices_sorted

unique_indices