jax.experimental.sparse.bcoo_broadcast_in_dim#

jax.experimental.sparse.bcoo_broadcast_in_dim(mat, *, shape, broadcast_dimensions, sharding=None)[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 (Shape) – 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.

Returns:

A BCOO-format array containing the target array.

Return type:

BCOO