JAX 中文文档（十五）（5）

JAX 中文文档（十五）（4）https://developer.aliyun.com/article/1559771

稀疏 API 参考

BCOO 数据结构

BCOO 是 Batched COO format，是在 jax.experimental.sparse 中实现的主要稀疏数据结构。其操作与 JAX 的核心转换兼容，包括批处理（例如 jax.vmap()）和自动微分（例如 jax.grad()）。

BCSR 数据结构

BCSR是批量压缩稀疏行格式，正在开发中。其操作与 JAX 的核心转换兼容，包括批处理（如jax.vmap()）和自动微分（如jax.grad()）。

其他稀疏数据结构

其他稀疏数据结构包括COO、CSR和CSC。这些是简单稀疏结构的参考实现，具有少数核心操作。它们的操作通常与自动微分转换（如jax.grad()）兼容，但不与批处理转换（如jax.vmap()）兼容。

jax.experimental.sparse.linalg

稀疏线性代数例程。

jax.experimental.sparse.BCOO

原文：jax.readthedocs.io/en/latest/_autosummary/jax.experimental.sparse.BCOO.html

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

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

参数：

• **(**data – 批量 COO 格式中的数据和索引。
  
• indices**)** – 批量 COO 格式中的数据和索引。
  
• shape (tuple[int, …**]) – 稀疏数组的形状。
  
• args (tuple[Array,* Array]*)
  
• indices_sorted (bool)
  
• unique_indices (bool)
  

data

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

类型：

jax.Array

indices

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

类型：

jax.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)

参数：

• args (tuple[Array,* Array]*)
  
• shape (Sequence[int])
  
• indices_sorted (bool)
  
• unique_indices (bool)
  

方法

属性

jax.experimental.sparse.bcoo_broadcast_in_dim

原文：jax.readthedocs.io/en/latest/_autosummary/jax.experimental.sparse.bcoo_broadcast_in_dim.html

jax.experimental.sparse.bcoo_broadcast_in_dim(mat, *, shape, broadcast_dimensions)

通过复制数据扩展 BCOO 数组的大小和秩。

BCOO 相当于 jax.lax.broadcast_in_dim。

参数：

• mat（BCOO） – BCOO 格式的数组。
  
• shape（tuple[int,* …]*） – 目标数组的形状。
  
• broadcast_dimensions（Sequence[int]） – 目标数组形状的维度，每个操作数（mat）形状对应一个维度。
  

返回：

包含目标数组的 BCOO 格式数组。

返回类型：

BCOO

jax.experimental.sparse.bcoo_concatenate

原文：jax.readthedocs.io/en/latest/_autosummary/jax.experimental.sparse.bcoo_concatenate.html

jax.experimental.sparse.bcoo_concatenate(operands, *, dimension)

稀疏实现的jax.lax.concatenate()函数

参数：

• operands（Sequence[BCOO]） – 要连接的 BCOO 数组序列。这些数组必须具有相同的形状，除了在维度轴上。此外，这些数组必须具有等效的批处理、稀疏和密集维度。
  
• dimension（int） – 指定沿其连接数组的维度的正整数。维度必须是输入的批处理或稀疏维度之一；不支持沿密集维度的连接。
  

返回值：

包含输入数组连接的 BCOO 数组。

返回类型：

BCOO

jax.experimental.sparse.bcoo_dot_general

原文：jax.readthedocs.io/en/latest/_autosummary/jax.experimental.sparse.bcoo_dot_general.html

jax.experimental.sparse.bcoo_dot_general(lhs, rhs, *, dimension_numbers, precision=None, preferred_element_type=None)

一般的收缩操作。

参数:

• lhs（BCOO | Array） – 一个 ndarray 或 BCOO 格式的稀疏数组。
  
• rhs（BCOO | Array） – 一个 ndarray 或 BCOO 格式的稀疏数组。
  
• dimension_numbers（tuple[tuple[Sequence[int]**, Sequence[int]], tuple[Sequence[int]**, Sequence[int]]]) – 一个形如((lhs_contracting_dims, rhs_contracting_dims), (lhs_batch_dims, rhs_batch_dims))的元组的元组。
  
• precision（None） – 未使用
  
• preferred_element_type（None） – 未使用
  

返回:

一个包含结果的 ndarray 或 BCOO 格式的稀疏数组。如果两个输入都是稀疏的，结果将是稀疏的，类型为 BCOO。如果任一输入是密集的，结果将是密集的，类型为 ndarray。

返回类型:

BCOO | Array

jax.experimental.sparse.bcoo_dot_general_sampled

原文：jax.readthedocs.io/en/latest/_autosummary/jax.experimental.sparse.bcoo_dot_general_sampled.html

jax.experimental.sparse.bcoo_dot_general_sampled(A, B, indices, *, dimension_numbers)

给定稀疏索引处计算输出的收缩操作。

参数：

• lhs – 一个 ndarray。
  
• rhs – 一个 ndarray。
  
• indices（Array） – BCOO 索引。
  
• dimension_numbers（tuple[tuple[Sequence[int]**, Sequence[int]], tuple[Sequence[int]**, Sequence[int]]]) – 形式为 ((lhs 收缩维度，rhs 收缩维度)，(lhs 批次维度，rhs 批次维度)) 的元组的元组。
  
• A（Array）
  
• B（Array）
  

返回：

BCOO 数据，包含结果的 ndarray。

返回类型：

Array

jax.experimental.sparse.bcoo_dynamic_slice

原文：jax.readthedocs.io/en/latest/_autosummary/jax.experimental.sparse.bcoo_dynamic_slice.html

jax.experimental.sparse.bcoo_dynamic_slice(mat, start_indices, slice_sizes)

{func}jax.lax.dynamic_slice的稀疏实现。

参数：

• mat (BCOO) – 要切片的 BCOO 数组。
  
• start_indices (Sequence[Any]) – 每个维度的标量索引列表。这些值可能是动态的。
  
• slice_sizes (Sequence[int]) – 切片的大小。必须是非负整数序列，长度等于操作数的维度数。在 JIT 编译的函数内部，仅支持静态值（所有 JAX 数组在 JIT 内必须具有静态已知大小）。
  

返回：

包含切片的 BCOO 数组。

返回类型：

out

jax.experimental.sparse.bcoo_extract

原文：jax.readthedocs.io/en/latest/_autosummary/jax.experimental.sparse.bcoo_extract.html

jax.experimental.sparse.bcoo_extract(sparr, arr, *, assume_unique=None)

根据稀疏数组的索引从密集数组中提取值。

参数：

• sparr (BCOO) – 用于输出的 BCOO 数组的索引。
  
• arr (jax.typing.ArrayLike) – 形状与 self.shape 相同的 ArrayLike
  
• assume_unique (bool | None) – 布尔值，默认为 sparr.unique_indices。如果为 True，则提取每个索引的值，即使索引包含重复项。如果为 False，则重复的索引将其值求和，并返回第一个索引的位置。
  

返回：

一个具有与 self 相同稀疏模式的 BCOO 数组。

返回类型：

提取的结果

jax.experimental.sparse.bcoo_fromdense

原文：jax.readthedocs.io/en/latest/_autosummary/jax.experimental.sparse.bcoo_fromdense.html

jax.experimental.sparse.bcoo_fromdense(mat, *, nse=None, n_batch=0, n_dense=0, index_dtype=<class 'jax.numpy.int32'>)

从密集矩阵创建 BCOO 格式的稀疏矩阵。

参数：

• mat（Array）– 要转换为 BCOO 格式的数组。
  
• nse（int | None）– 每个批次中指定元素的数量
  
• n_batch（int）– 批次维度的数量（默认：0）
  
• n_dense（int）– 块维度的数量（默认：0）
  
• index_dtype（jax.typing.DTypeLike）– 稀疏索引的数据类型（默认：int32）
  

返回：

矩阵的 BCOO 表示。

返回类型：

mat_bcoo

jax.experimental.sparse.bcoo_gather

原文：jax.readthedocs.io/en/latest/_autosummary/jax.experimental.sparse.bcoo_gather.html

jax.experimental.sparse.bcoo_gather(operand, start_indices, dimension_numbers, slice_sizes, *, unique_indices=False, indices_are_sorted=False, mode=None, fill_value=None)

BCOO 版本的 lax.gather。

参数：

• operand (BCOO)
  
• start_indices (数组)
  
• dimension_numbers (GatherDimensionNumbers)
  
• slice_sizes (tuple[int, …**])
  
• unique_indices (bool)
  
• indices_are_sorted (bool)
  
• mode (str | GatherScatterMode | None)
  

返回类型：

BCOO