Getting Started
Further Resources
jax.debug.print
jax.debug.breakpoint
checkify
xmap
jax.numpy
jax.scipy
jax.lax
jax.random
jax.sharding
jax.debug
jax.dlpack
jax.distributed
jax.flatten_util
jax.image
jax.nn
jax.nn.initializers
jax.ops
jax.profiler
jax.stages
jax.tree_util
jax.example_libraries
jax.example_libraries.optimizers
jax.example_libraries.stax
jax.experimental
jax.experimental.checkify
jax.experimental.global_device_array
jax.experimental.host_callback
jax.experimental.maps
jax.experimental.pjit
jax.experimental.sparse
jax.experimental.jet
jax.experimental.custom_partitioning
jax.lib
Expand the size and rank of a BCOO array by duplicating the data.
A BCOO equivalence to jax.lax.broadcast_in_dim.
mat (BCOO) – A BCOO-format array.
BCOO
shape (Tuple[int, ...]) – The shape of the target array.
Tuple
int
...
broadcast_dimensions (Sequence[int]) – The dimension in the shape of the target array which each dimension of the operand (mat) shape corresponds to.
Sequence
mat
A BCOO-format array containing the target array.