def gather_fn(x):
return partial_transpose(
gather_flat(
partial_transpose(x, [1, 0]),
batch_beam_ids,
self.batch_size,
self.beam_size),
[1, 0])
def gather_fn(x):
if len(x.shape.dims) < 2:
return x
return partial_transpose(
gather_flat(
partial_transpose(x, [1, 0]),
batch_beam_ids,
self.batch_size,
self.beam_size),
[1, 0])