"""Python wrappers around TensorFlow ops.

This file is MACHINE GENERATED! Do not edit.
Original C++ source file: beam_search_ops.cc
"""

import collections as _collections

from tensorflow.python.eager import execute as _execute
from tensorflow.python.eager import context as _context
from tensorflow.python.eager import core as _core
from tensorflow.python.framework import dtypes as _dtypes
from tensorflow.python.framework import tensor_shape as _tensor_shape

from tensorflow.core.framework import op_def_pb2 as _op_def_pb2
# Needed to trigger the call to _set_call_cpp_shape_fn.
from tensorflow.python.framework import common_shapes as _common_shapes
from tensorflow.python.framework import op_def_registry as _op_def_registry
from tensorflow.python.framework import ops as _ops
from tensorflow.python.framework import op_def_library as _op_def_library
from tensorflow.python.util.tf_export import tf_export


@tf_export('gather_tree')
def gather_tree(step_ids, parent_ids, max_sequence_lengths, end_token, name=None):
  r"""Calculates the full beams from the per-step ids and parent beam ids.

  On CPU, if an out of bound parent id is found, an error is returned.
  On GPU, if an out of bound parent id is found, a -1 is stored in the
  corresponding output value and the execution for that beam returns early.

  For a given beam, past the time step containing the first decoded `end_token`
  all values are filled in with `end_token`.

  TODO(ebrevdo): fill in the remainder of this docstring.

  Args:
    step_ids: A `Tensor`. Must be one of the following types: `int32`.
      `[max_time, batch_size, beam_width]`.
    parent_ids: A `Tensor`. Must have the same type as `step_ids`.
      `[max_time, batch_size, beam_width]`.
    max_sequence_lengths: A `Tensor` of type `int32`. `[batch_size]`.
    end_token: A `Tensor`. Must have the same type as `step_ids`. `[]`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor`. Has the same type as `step_ids`.
    `[max_time, batch_size, beam_width]`.
  """
  _ctx = _context.context()
  if _ctx.in_graph_mode():
    _, _, _op = _op_def_lib._apply_op_helper(
        "GatherTree", step_ids=step_ids, parent_ids=parent_ids,
        max_sequence_lengths=max_sequence_lengths, end_token=end_token,
        name=name)
    _result = _op.outputs[:]
    _inputs_flat = _op.inputs
    _attrs = ("T", _op.get_attr("T"))
  else:
    _attr_T, _inputs_T = _execute.args_to_matching_eager([step_ids, parent_ids, end_token], _ctx)
    (step_ids, parent_ids, end_token) = _inputs_T
    max_sequence_lengths = _ops.convert_to_tensor(max_sequence_lengths, _dtypes.int32)
    _inputs_flat = [step_ids, parent_ids, max_sequence_lengths, end_token]
    _attrs = ("T", _attr_T)
    _result = _execute.execute(b"GatherTree", 1, inputs=_inputs_flat,
                               attrs=_attrs, ctx=_ctx, name=name)
  _execute.record_gradient(
      "GatherTree", _inputs_flat, _attrs, _result, name)
  _result, = _result
  return _result

_ops.RegisterShape("GatherTree")(None)

def _InitOpDefLibrary(op_list_proto_bytes):
  op_list = _op_def_pb2.OpList()
  op_list.ParseFromString(op_list_proto_bytes)
  _op_def_registry.register_op_list(op_list)
  op_def_lib = _op_def_library.OpDefLibrary()
  op_def_lib.add_op_list(op_list)
  return op_def_lib
# op {
#   name: "GatherTree"
#   input_arg {
#     name: "step_ids"
#     type_attr: "T"
#   }
#   input_arg {
#     name: "parent_ids"
#     type_attr: "T"
#   }
#   input_arg {
#     name: "max_sequence_lengths"
#     type: DT_INT32
#   }
#   input_arg {
#     name: "end_token"
#     type_attr: "T"
#   }
#   output_arg {
#     name: "beams"
#     type_attr: "T"
#   }
#   attr {
#     name: "T"
#     type: "type"
#     allowed_values {
#       list {
#         type: DT_INT32
#       }
#     }
#   }
# }
_op_def_lib = _InitOpDefLibrary(b"\nt\n\nGatherTree\022\r\n\010step_ids\"\001T\022\017\n\nparent_ids\"\001T\022\030\n\024max_sequence_lengths\030\003\022\016\n\tend_token\"\001T\032\n\n\005beams\"\001T\"\020\n\001T\022\004type:\005\n\0032\001\003")
