jax.experimental.sparse.bcoo_concatenate

jax.experimental.sparse.bcoo_concatenate#

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

Sparse implementation of jax.lax.concatenate()

Parameters:
  • operands (Sequence[BCOO]) – Sequence of BCOO arrays to concatenate. The arrays must have equal shapes, except in the dimension axis. Additionally, the arrays must have have equivalent batch, sparse, and dense dimensions.

  • dimension (int) – Positive integer specifying the dimension along which to concatenate the arrays. The dimension must be among batch or sparse dimensions of the input; concatenation along dense dimensions is not supported.

Return type:

BCOO

Returns:

A BCOO array containing the concatenation of the inputs.