jax.experimental.sparse.bcoo_concatenate

目录

jax.experimental.sparse.bcoo_concatenate#

jax.experimental.sparse.bcoo_concatenate(operands, *, dimension)[源码]#

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

参数:

operands (Sequence[BCOO]) – 要连接的 BCOO 数组序列。数组在 dimension 轴之外的形状必须相等。此外，数组必须具有相同的批次、稀疏和密集维度。
dimension (int) – 指定沿哪个维度连接数组的正整数。该维度必须是输入批次或稀疏维度之一；不支持沿密集维度连接。

返回:

包含输入连接的 BCOO 数组。

返回类型: