jax.experimental.sparse.bcoo_broadcast_in_dim#

jax.experimental.sparse.bcoo_broadcast_in_dim(mat, *, shape, broadcast_dimensions)[source]#

Expand the size and rank of a BCOO array by duplicating the data.

A BCOO equivalence to jax.lax.broadcast_in_dim.

Parameters

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

  • shape (tuple[int, ...]) – The shape of the target array.

  • broadcast_dimensions (Sequence[int]) – The dimension in the shape of the target array which each dimension of the operand (mat) shape corresponds to.

Return type

BCOO

Returns

A BCOO-format array containing the target array.