jax.experimental.sparse.bcoo_broadcast_in_dim

jax.experimental.sparse.bcoo_broadcast_in_dim(mat, *, shape, broadcast_dimensions)

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.