jax.experimental.sparse.bcoo_transpose#

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

Transpose a BCOO-format array.

Parameters
  • mat – 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.

Returns

A BCOO-format array.