jax.experimental.sparse.bcoo_gather

Contents

jax.experimental.sparse.bcoo_gather#

jax.experimental.sparse.bcoo_gather(operand, start_indices, dimension_numbers, slice_sizes, *, unique_indices=False, indices_are_sorted=False, mode=None, fill_value=None)[source]#

BCOO version of lax.gather.

Parameters:
  • operand (BCOO)

  • start_indices (Array)

  • dimension_numbers (GatherDimensionNumbers)

  • slice_sizes (Shape)

  • unique_indices (bool)

  • indices_are_sorted (bool)

  • mode (str | GatherScatterMode | None)

Return type:

BCOO