def expand_to_beam(self, val: tf.Tensor, dim: int = 0) -> tf.Tensor:
"""Copy a tensor along a new beam dimension.
Arguments:
val: The ``Tensor`` to expand.
dim: The dimension along which to expand. Usually, the batch axis.
Returns:
The expanded tensor.
"""
orig_shape = get_shape_list(val)
if val.shape.ndims == 0:
return val
orig_shape[dim] *= self.beam_size
tile_shape = [1] * (len(orig_shape) + 1)
tile_shape[dim + 1] = self.beam_size
val = tf.tile(tf.expand_dims(val, 1), tile_shape)
val = tf.reshape(val, orig_shape)
return val
def expand_to_beam(self, val: tf.Tensor, dim: int = 0) -> tf.Tensor:
"""Copy a tensor along a new beam dimension.
Arguments:
val: The ``Tensor`` to expand.
dim: The dimension along which to expand. Usually, the batch axis.
Returns:
The expanded tensor.
"""
orig_shape = get_shape_list(val)
if val.shape.ndims == 0:
return val
orig_shape[dim] *= self.beam_size
tile_shape = [1] * (len(orig_shape) + 1)
tile_shape[dim + 1] = self.beam_size
val = tf.tile(tf.expand_dims(val, 1), tile_shape)
val = tf.reshape(val, orig_shape)
return val