jax.experimental.sparse.BCOO#

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

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)[原始碼]#
參數:

方法

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

astype(*args, **kwargs)

複製陣列並轉換為指定的 dtype。

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