MultiBatchBeam.multi_batch_beam(array, start_idxs, batch_lens, beam_width, wrap_mode, pad_left=0, pad_right=0, idx_dim=0, batch_dim=1)[source]
  • array – ndarray, at least 2D. symbolic
  • start_idxs – ndarray, 1D. symbolic. can be float (for gpu)
  • batch_lens – ndarray, 1D. symbolic. len of each batch. can be float (for gpu)
  • beam_width – scalar. symbolic.
  • wrap_mode – “wrap_around” or “pad”. static.
  • idx_dim – int. where to apply each start_idxs[i]. static.
  • batch_dim – the same dim as in start_idxs. static.
  • pad_value – used in wrap_mode “pad”. automatically broadcasted. symbolic.

ndarray like array, but shape[idx_dim] == beam_width

See also _naive_multi_batch_beam for one naive reference implementation.

class MultiBatchBeam.MultiBatchBeamOp(wrap_mode, idx_dim=0, batch_dim=1)[source]
make_node(array, start_idxs, batch_lens, beam_width, pad_left, pad_right)[source]
perform(node, inputs, output_storage)[source]
infer_shape(node, input_shapes)[source]
grad(inputs, output_grads)[source]
class MultiBatchBeam.MultiBatchBeamGradAddOp(wrap_mode, idx_dim=0, batch_dim=1, inplace=False, zero_with_shape=False, array_ndim=None)[source]

(D_array / D_array_shape, start_idxs, batch_lens, beam_width, D_beam) -> D_array + grad

  • wrap_mode (str) – “wrap_around” or “pad”
  • idx_dim (int) – usually that’s time dim
  • batch_dim (int) – batch dim
  • inplace (bool) – operate inplace on input
  • zero_with_shape (bool) – we get D_array_shape as the first input and init D_array with zero
  • array_ndim (int) – ndim of array/D_array. needed for zero_with_shape
make_node(D_array_or_shape, start_idxs, batch_lens, beam_width, D_beam)[source]
infer_shape(node, input_shapes)[source]
perform(node, inputs, output_storage)[source]
class MultiBatchBeam.GpuMultiBatchBeamOp[source]