Function gather_fn

axlearn/common/decoding.py:108–109

(x)

        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.