jax.experimental.sparse.bcoo_transpose

jax.experimental.sparse.bcoo_transpose(mat, *, permutation)

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 permutating dense axes with non-dense axes.

Return type:

BCOO

Returns:

A BCOO-format array.