jax.experimental.sparse.bcoo_transpose

Contents

jax.experimental.sparse.bcoo_transpose#

jax.experimental.sparse.bcoo_transpose(mat, *, permutation)[source]#

Transpose a BCOO-format array.

Parameters:

  • mat (BCOO) – A BCOO-format array.

  • permutation (Sequence[int]) – A tuple or list or ndarray which contains a permutation of [0,1,..,N-1] where N is the number of axes of mat in the order of batch, sparse, and dense dimensions. The i’th axis of the returned array corresponds to the axis numbered permutation[i] of mat. Transpose permutation currently does not support permuting batch axes with non-batch axes nor permuting dense axes with non-dense axes.

Returns:

A BCOO-format array.

Return type:

BCOO

Contents