Function _gather_beams

axlearn/common/decoding.py:78–120

Gathers the beam slices indexed by beam_indices into new beam array. Args: nested: NestedTensor or scalars (the latter ignored). beam_indices: Array of beam_indices. batch_size: Size of batch. old_beam_size: Size of _old_ beam dimension. new_beam_size: Si

(
    nested: NestedTensor,
    beam_indices: Tensor,
    batch_size: int,
    old_beam_size: int,
    new_beam_size: int,
    one_hot: bool = True,
)

def _gather_beams(
    nested: NestedTensor,
    beam_indices: Tensor,
    batch_size: int,
    old_beam_size: int,
    new_beam_size: int,
    one_hot: bool = True,
) -> Tensor:
    """Gathers the beam slices indexed by beam_indices into new beam array.

    Args:
        nested: NestedTensor or scalars (the latter ignored).
        beam_indices: Array of beam_indices.
        batch_size: Size of batch.
        old_beam_size: Size of _old_ beam dimension.
        new_beam_size: Size of _new_ beam dimension.
        one_hot: Whether to perform gathers by one-hot contraction or directly.

    Returns:
        New beam arrays
        [batch_size, old_beam_size, ...] --> [batch_size, new_beam_size, ...].

    TODO(sneha): The method does a likely O(n) copy to track history,
          could use a tree / single parent pointer instead.
          Check with timeit after.
    """
    if one_hot:
        # Gather via one-hot contraction, needed for SPMD partitioning.
        oh_beam_indices = jax.nn.one_hot(beam_indices, old_beam_size, dtype=jnp.int32)

        def gather_fn(x):
            return jnp.einsum("bno,bo...->bn...", oh_beam_indices, x).astype(x.dtype)

    else:
        # True gather via fancy indexing.
        batch_indices = jnp.reshape(
            jnp.arange(batch_size * new_beam_size) // new_beam_size, (batch_size, new_beam_size)
        )

        def gather_fn(x):
            return x[batch_indices, beam_indices]

    return vectorized_tree_map(lambda x: gather_fn(x) if x.ndim > 1 else x, nested)


def _top_k_two_stage(x: Tensor, k: int):