"""Python wrappers around TensorFlow ops.

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

import collections

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

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.deprecation import deprecated_endpoints
from tensorflow.python.util import dispatch as _dispatch
from tensorflow.python.util.tf_export import tf_export

from typing import TypeVar
_BatchOutput = collections.namedtuple(
    "Batch",
    ["batched_tensors", "batch_index", "id"])


def batch(in_tensors, num_batch_threads, max_batch_size, batch_timeout_micros, grad_timeout_micros, max_enqueued_batches=10, allowed_batch_sizes=[], container="", shared_name="", batching_queue="", name=None):
  r"""Batches all input tensors nondeterministically.

  When many instances of this Op are being run concurrently with the same
  container/shared_name in the same device, some will output zero-shaped Tensors
  and others will output Tensors of size up to max_batch_size.

  All Tensors in in_tensors are batched together (so, for example, labels and
  features should be batched with a single instance of this operation.

  Each invocation of batch emits an `id` scalar which will be used to identify
  this particular invocation when doing unbatch or its gradient.

  Each op which emits a non-empty batch will also emit a non-empty batch_index
  Tensor, which, is a [K, 3] matrix where each row contains the invocation's id,
  start, and length of elements of each set of Tensors present in batched_tensors.

  Batched tensors are concatenated along the first dimension, and all tensors in
  in_tensors must have the first dimension of the same size.

  in_tensors: The tensors to be batched.
  num_batch_threads: Number of scheduling threads for processing batches of work.
   Determines the number of batches processed in parallel.
  max_batch_size: Batch sizes will never be bigger than this.
  batch_timeout_micros: Maximum number of microseconds to wait before outputting
   an incomplete batch.
  allowed_batch_sizes: Optional list of allowed batch sizes. If left empty, does
   nothing. Otherwise, supplies a list of batch sizes, causing the op to pad
   batches up to one of those sizes. The entries must increase monotonically, and
   the final entry must equal max_batch_size.
  grad_timeout_micros: The timeout to use for the gradient. See Unbatch.
  batched_tensors: Either empty tensors or a batch of concatenated Tensors.
  batch_index: If out_tensors is non-empty, has information to invert it.
  container: Controls the scope of sharing of this batch.
  id: always contains a scalar with a unique ID for this invocation of Batch.
  shared_name: Concurrently running instances of batch in the same device with the
   same container and shared_name will batch their elements together. If left
   empty, the op name will be used as the shared name.
  T: the types of tensors to be batched.

  Args:
    in_tensors: A list of `Tensor` objects.
    num_batch_threads: An `int`.
    max_batch_size: An `int`.
    batch_timeout_micros: An `int`.
    grad_timeout_micros: An `int`.
    max_enqueued_batches: An optional `int`. Defaults to `10`.
    allowed_batch_sizes: An optional list of `ints`. Defaults to `[]`.
    container: An optional `string`. Defaults to `""`.
    shared_name: An optional `string`. Defaults to `""`.
    batching_queue: An optional `string`. Defaults to `""`.
    name: A name for the operation (optional).

  Returns:
    A tuple of `Tensor` objects (batched_tensors, batch_index, id).

    batched_tensors: A list of `Tensor` objects. Has the same type as `in_tensors`.
    batch_index: A `Tensor` of type `int64`.
    id: A `Tensor` of type `int64`.
  """
  _ctx = _context._context or _context.context()
  tld = _ctx._thread_local_data
  if tld.is_eager:
    try:
      _result = pywrap_tfe.TFE_Py_FastPathExecute(
        _ctx, "Batch", name, in_tensors, "num_batch_threads",
        num_batch_threads, "max_batch_size", max_batch_size,
        "max_enqueued_batches", max_enqueued_batches, "batch_timeout_micros",
        batch_timeout_micros, "allowed_batch_sizes", allowed_batch_sizes,
        "grad_timeout_micros", grad_timeout_micros, "container", container,
        "shared_name", shared_name, "batching_queue", batching_queue)
      _result = _BatchOutput._make(_result)
      return _result
    except _core._NotOkStatusException as e:
      _ops.raise_from_not_ok_status(e, name)
    except _core._FallbackException:
      pass
    try:
      return batch_eager_fallback(
          in_tensors, num_batch_threads=num_batch_threads,
          max_batch_size=max_batch_size,
          max_enqueued_batches=max_enqueued_batches,
          batch_timeout_micros=batch_timeout_micros,
          allowed_batch_sizes=allowed_batch_sizes,
          grad_timeout_micros=grad_timeout_micros, container=container,
          shared_name=shared_name, batching_queue=batching_queue, name=name,
          ctx=_ctx)
    except _core._SymbolicException:
      pass  # Add nodes to the TensorFlow graph.
  # Add nodes to the TensorFlow graph.
  num_batch_threads = _execute.make_int(num_batch_threads, "num_batch_threads")
  max_batch_size = _execute.make_int(max_batch_size, "max_batch_size")
  batch_timeout_micros = _execute.make_int(batch_timeout_micros, "batch_timeout_micros")
  grad_timeout_micros = _execute.make_int(grad_timeout_micros, "grad_timeout_micros")
  if max_enqueued_batches is None:
    max_enqueued_batches = 10
  max_enqueued_batches = _execute.make_int(max_enqueued_batches, "max_enqueued_batches")
  if allowed_batch_sizes is None:
    allowed_batch_sizes = []
  if not isinstance(allowed_batch_sizes, (list, tuple)):
    raise TypeError(
        "Expected list for 'allowed_batch_sizes' argument to "
        "'batch' Op, not %r." % allowed_batch_sizes)
  allowed_batch_sizes = [_execute.make_int(_i, "allowed_batch_sizes") for _i in allowed_batch_sizes]
  if container is None:
    container = ""
  container = _execute.make_str(container, "container")
  if shared_name is None:
    shared_name = ""
  shared_name = _execute.make_str(shared_name, "shared_name")
  if batching_queue is None:
    batching_queue = ""
  batching_queue = _execute.make_str(batching_queue, "batching_queue")
  _, _, _op, _outputs = _op_def_library._apply_op_helper(
        "Batch", in_tensors=in_tensors, num_batch_threads=num_batch_threads,
                 max_batch_size=max_batch_size,
                 batch_timeout_micros=batch_timeout_micros,
                 grad_timeout_micros=grad_timeout_micros,
                 max_enqueued_batches=max_enqueued_batches,
                 allowed_batch_sizes=allowed_batch_sizes, container=container,
                 shared_name=shared_name, batching_queue=batching_queue,
                 name=name)
  _result = _outputs[:]
  if _execute.must_record_gradient():
    _attrs = ("num_batch_threads", _op._get_attr_int("num_batch_threads"),
              "max_batch_size", _op._get_attr_int("max_batch_size"),
              "max_enqueued_batches",
              _op._get_attr_int("max_enqueued_batches"),
              "batch_timeout_micros",
              _op._get_attr_int("batch_timeout_micros"),
              "allowed_batch_sizes", _op.get_attr("allowed_batch_sizes"),
              "grad_timeout_micros", _op._get_attr_int("grad_timeout_micros"),
              "container", _op.get_attr("container"), "shared_name",
              _op.get_attr("shared_name"), "batching_queue",
              _op.get_attr("batching_queue"), "T", _op.get_attr("T"))
    _inputs_flat = _op.inputs
    _execute.record_gradient(
        "Batch", _inputs_flat, _attrs, _result)
  _result = [_result[:len(in_tensors)]] + _result[len(in_tensors):]
  _result = _BatchOutput._make(_result)
  return _result

Batch = tf_export("raw_ops.Batch")(_ops.to_raw_op(batch))


def batch_eager_fallback(in_tensors, num_batch_threads, max_batch_size, batch_timeout_micros, grad_timeout_micros, max_enqueued_batches, allowed_batch_sizes, container, shared_name, batching_queue, name, ctx):
  num_batch_threads = _execute.make_int(num_batch_threads, "num_batch_threads")
  max_batch_size = _execute.make_int(max_batch_size, "max_batch_size")
  batch_timeout_micros = _execute.make_int(batch_timeout_micros, "batch_timeout_micros")
  grad_timeout_micros = _execute.make_int(grad_timeout_micros, "grad_timeout_micros")
  if max_enqueued_batches is None:
    max_enqueued_batches = 10
  max_enqueued_batches = _execute.make_int(max_enqueued_batches, "max_enqueued_batches")
  if allowed_batch_sizes is None:
    allowed_batch_sizes = []
  if not isinstance(allowed_batch_sizes, (list, tuple)):
    raise TypeError(
        "Expected list for 'allowed_batch_sizes' argument to "
        "'batch' Op, not %r." % allowed_batch_sizes)
  allowed_batch_sizes = [_execute.make_int(_i, "allowed_batch_sizes") for _i in allowed_batch_sizes]
  if container is None:
    container = ""
  container = _execute.make_str(container, "container")
  if shared_name is None:
    shared_name = ""
  shared_name = _execute.make_str(shared_name, "shared_name")
  if batching_queue is None:
    batching_queue = ""
  batching_queue = _execute.make_str(batching_queue, "batching_queue")
  _attr_T, in_tensors = _execute.convert_to_mixed_eager_tensors(in_tensors, ctx)
  _inputs_flat = list(in_tensors)
  _attrs = ("num_batch_threads", num_batch_threads, "max_batch_size",
  max_batch_size, "max_enqueued_batches", max_enqueued_batches,
  "batch_timeout_micros", batch_timeout_micros, "allowed_batch_sizes",
  allowed_batch_sizes, "grad_timeout_micros", grad_timeout_micros,
  "container", container, "shared_name", shared_name, "batching_queue",
  batching_queue, "T", _attr_T)
  _result = _execute.execute(b"Batch", len(in_tensors) + 2,
                             inputs=_inputs_flat, attrs=_attrs, ctx=ctx,
                             name=name)
  if _execute.must_record_gradient():
    _execute.record_gradient(
        "Batch", _inputs_flat, _attrs, _result)
  _result = [_result[:len(in_tensors)]] + _result[len(in_tensors):]
  _result = _BatchOutput._make(_result)
  return _result


def batch_function(in_tensors, captured_tensors, f, num_batch_threads, max_batch_size, batch_timeout_micros, Tout, max_enqueued_batches=10, allowed_batch_sizes=[], container="", shared_name="", batching_queue="", enable_large_batch_splitting=False, name=None):
  r"""Batches all the inputs tensors to the computation done by the function.

  So, for example, in the following code

    ```python

    # This input will be captured.
    y = tf.placeholder_with_default(1.0, shape=[])

    @tf.Defun(tf.float32)
    def computation(a):
      return tf.matmul(a, a) + y

    b = gen_batch_ops.batch_function(
            f=computation
            in_tensors=[a],
            captured_tensors=computation.captured_inputs,
            Tout=[o.type for o in computation.definition.signature.output_arg],
            num_batch_threads=1,
            max_batch_size=10,
            batch_timeout_micros=100000,  # 100ms
            allowed_batch_sizes=[3, 10],
            batching_queue="")
    ```

  If more than one session.run call is simultaneously trying to compute `b`
  the values of `a` will be gathered, non-deterministically concatenated
  along the first axis, and only one thread will run the computation.

  Assumes that all arguments of the function are Tensors which will be batched
  along their first dimension.

  Arguments that are captured, are not batched. The session.run call which does
  the concatenation, will use the values of the captured tensors available to it.
  Therefore, typical uses of captured tensors should involve values which remain
  unchanged across session.run calls. Inference is a good example of this.

  SparseTensor is not supported. The return value of the decorated function
  must be a Tensor or a list/tuple of Tensors.

  Args:
    in_tensors: A list of `Tensor` objects. The tensors to be batched.
    captured_tensors: A list of `Tensor` objects.
      The tensors which are captured in the function, and don't need
      to be batched.
    f: A function decorated with @Defun.
    num_batch_threads: An `int`.
      Number of scheduling threads for processing batches of work.
      Determines the number of batches processed in parallel.
    max_batch_size: An `int`. Batch sizes will never be bigger than this.
    batch_timeout_micros: An `int`.
      Maximum number of microseconds to wait before outputting
      an incomplete batch.
    Tout: A list of `tf.DTypes` that has length `>= 1`.
      the types of the output tensors.
    max_enqueued_batches: An optional `int`. Defaults to `10`.
      Maximum number of batches enqueued. Default: 10.
    allowed_batch_sizes: An optional list of `ints`. Defaults to `[]`.
      Optional list of allowed batch sizes. If left empty, does
      nothing. Otherwise, supplies a list of batch sizes, causing the op to pad
      batches up to one of those sizes. The entries must increase monotonically.
      If enable_large_batch_splitting is false (i.e., large-input-split is not
      enabled) the final entry must equal max_batch_size.
    container: An optional `string`. Defaults to `""`.
      Controls the scope of sharing of this batch.
    shared_name: An optional `string`. Defaults to `""`.
      Concurrently running instances of batch in the same device with the
      same container and shared_name will batch their elements together. If left
      empty, the op name will be used as the shared name.
    batching_queue: An optional `string`. Defaults to `""`.
    enable_large_batch_splitting: An optional `bool`. Defaults to `False`.
      input with a large size (i.e., larger than the largest value of
      `allowed_batch_sizes`) will be splitted into multiple batches with batch size.
    name: A name for the operation (optional).

  Returns:
    A list of `Tensor` objects of type `Tout`.
  """
  _ctx = _context._context or _context.context()
  tld = _ctx._thread_local_data
  if tld.is_eager:
    try:
      _result = pywrap_tfe.TFE_Py_FastPathExecute(
        _ctx, "BatchFunction", name, in_tensors, captured_tensors, "f", f,
        "num_batch_threads", num_batch_threads, "max_batch_size",
        max_batch_size, "batch_timeout_micros", batch_timeout_micros,
        "max_enqueued_batches", max_enqueued_batches, "allowed_batch_sizes",
        allowed_batch_sizes, "container", container, "shared_name",
        shared_name, "batching_queue", batching_queue, "Tout", Tout,
        "enable_large_batch_splitting", enable_large_batch_splitting)
      return _result
    except _core._NotOkStatusException as e:
      _ops.raise_from_not_ok_status(e, name)
    except _core._FallbackException:
      pass
    try:
      return batch_function_eager_fallback(
          in_tensors, captured_tensors, f=f,
          num_batch_threads=num_batch_threads, max_batch_size=max_batch_size,
          batch_timeout_micros=batch_timeout_micros,
          max_enqueued_batches=max_enqueued_batches,
          allowed_batch_sizes=allowed_batch_sizes, container=container,
          shared_name=shared_name, batching_queue=batching_queue, Tout=Tout,
          enable_large_batch_splitting=enable_large_batch_splitting,
          name=name, ctx=_ctx)
    except _core._SymbolicException:
      pass  # Add nodes to the TensorFlow graph.
  # Add nodes to the TensorFlow graph.
  num_batch_threads = _execute.make_int(num_batch_threads, "num_batch_threads")
  max_batch_size = _execute.make_int(max_batch_size, "max_batch_size")
  batch_timeout_micros = _execute.make_int(batch_timeout_micros, "batch_timeout_micros")
  if not isinstance(Tout, (list, tuple)):
    raise TypeError(
        "Expected list for 'Tout' argument to "
        "'batch_function' Op, not %r." % Tout)
  Tout = [_execute.make_type(_t, "Tout") for _t in Tout]
  if max_enqueued_batches is None:
    max_enqueued_batches = 10
  max_enqueued_batches = _execute.make_int(max_enqueued_batches, "max_enqueued_batches")
  if allowed_batch_sizes is None:
    allowed_batch_sizes = []
  if not isinstance(allowed_batch_sizes, (list, tuple)):
    raise TypeError(
        "Expected list for 'allowed_batch_sizes' argument to "
        "'batch_function' Op, not %r." % allowed_batch_sizes)
  allowed_batch_sizes = [_execute.make_int(_i, "allowed_batch_sizes") for _i in allowed_batch_sizes]
  if container is None:
    container = ""
  container = _execute.make_str(container, "container")
  if shared_name is None:
    shared_name = ""
  shared_name = _execute.make_str(shared_name, "shared_name")
  if batching_queue is None:
    batching_queue = ""
  batching_queue = _execute.make_str(batching_queue, "batching_queue")
  if enable_large_batch_splitting is None:
    enable_large_batch_splitting = False
  enable_large_batch_splitting = _execute.make_bool(enable_large_batch_splitting, "enable_large_batch_splitting")
  _, _, _op, _outputs = _op_def_library._apply_op_helper(
        "BatchFunction", in_tensors=in_tensors,
                         captured_tensors=captured_tensors, f=f,
                         num_batch_threads=num_batch_threads,
                         max_batch_size=max_batch_size,
                         batch_timeout_micros=batch_timeout_micros, Tout=Tout,
                         max_enqueued_batches=max_enqueued_batches,
                         allowed_batch_sizes=allowed_batch_sizes,
                         container=container, shared_name=shared_name,
                         batching_queue=batching_queue,
                         enable_large_batch_splitting=enable_large_batch_splitting,
                         name=name)
  _result = _outputs[:]
  if _execute.must_record_gradient():
    _attrs = ("f", _op.get_attr("f"), "num_batch_threads",
              _op._get_attr_int("num_batch_threads"), "max_batch_size",
              _op._get_attr_int("max_batch_size"), "batch_timeout_micros",
              _op._get_attr_int("batch_timeout_micros"),
              "max_enqueued_batches",
              _op._get_attr_int("max_enqueued_batches"),
              "allowed_batch_sizes", _op.get_attr("allowed_batch_sizes"),
              "container", _op.get_attr("container"), "shared_name",
              _op.get_attr("shared_name"), "batching_queue",
              _op.get_attr("batching_queue"), "Tin", _op.get_attr("Tin"),
              "Tcaptured", _op.get_attr("Tcaptured"), "Tout",
              _op.get_attr("Tout"), "enable_large_batch_splitting",
              _op._get_attr_bool("enable_large_batch_splitting"))
    _inputs_flat = _op.inputs
    _execute.record_gradient(
        "BatchFunction", _inputs_flat, _attrs, _result)
  return _result

BatchFunction = tf_export("raw_ops.BatchFunction")(_ops.to_raw_op(batch_function))


def batch_function_eager_fallback(in_tensors, captured_tensors, f, num_batch_threads, max_batch_size, batch_timeout_micros, Tout, max_enqueued_batches, allowed_batch_sizes, container, shared_name, batching_queue, enable_large_batch_splitting, name, ctx):
  num_batch_threads = _execute.make_int(num_batch_threads, "num_batch_threads")
  max_batch_size = _execute.make_int(max_batch_size, "max_batch_size")
  batch_timeout_micros = _execute.make_int(batch_timeout_micros, "batch_timeout_micros")
  if not isinstance(Tout, (list, tuple)):
    raise TypeError(
        "Expected list for 'Tout' argument to "
        "'batch_function' Op, not %r." % Tout)
  Tout = [_execute.make_type(_t, "Tout") for _t in Tout]
  if max_enqueued_batches is None:
    max_enqueued_batches = 10
  max_enqueued_batches = _execute.make_int(max_enqueued_batches, "max_enqueued_batches")
  if allowed_batch_sizes is None:
    allowed_batch_sizes = []
  if not isinstance(allowed_batch_sizes, (list, tuple)):
    raise TypeError(
        "Expected list for 'allowed_batch_sizes' argument to "
        "'batch_function' Op, not %r." % allowed_batch_sizes)
  allowed_batch_sizes = [_execute.make_int(_i, "allowed_batch_sizes") for _i in allowed_batch_sizes]
  if container is None:
    container = ""
  container = _execute.make_str(container, "container")
  if shared_name is None:
    shared_name = ""
  shared_name = _execute.make_str(shared_name, "shared_name")
  if batching_queue is None:
    batching_queue = ""
  batching_queue = _execute.make_str(batching_queue, "batching_queue")
  if enable_large_batch_splitting is None:
    enable_large_batch_splitting = False
  enable_large_batch_splitting = _execute.make_bool(enable_large_batch_splitting, "enable_large_batch_splitting")
  _attr_Tin, in_tensors = _execute.convert_to_mixed_eager_tensors(in_tensors, ctx)
  _attr_Tcaptured, captured_tensors = _execute.convert_to_mixed_eager_tensors(captured_tensors, ctx)
  _inputs_flat = list(in_tensors) + list(captured_tensors)
  _attrs = ("f", f, "num_batch_threads", num_batch_threads, "max_batch_size",
  max_batch_size, "batch_timeout_micros", batch_timeout_micros,
  "max_enqueued_batches", max_enqueued_batches, "allowed_batch_sizes",
  allowed_batch_sizes, "container", container, "shared_name", shared_name,
  "batching_queue", batching_queue, "Tin", _attr_Tin, "Tcaptured",
  _attr_Tcaptured, "Tout", Tout, "enable_large_batch_splitting",
  enable_large_batch_splitting)
  _result = _execute.execute(b"BatchFunction", len(Tout), inputs=_inputs_flat,
                             attrs=_attrs, ctx=ctx, name=name)
  if _execute.must_record_gradient():
    _execute.record_gradient(
        "BatchFunction", _inputs_flat, _attrs, _result)
  return _result


def unbatch(batched_tensor, batch_index, id, timeout_micros, container="", shared_name="", name=None):
  r"""Reverses the operation of Batch for a single output Tensor.

  An instance of Unbatch either receives an empty batched_tensor, in which case it
  asynchronously waits until the values become available from a concurrently
  running instance of Unbatch with the same container and shared_name, or receives
  a non-empty batched_tensor in which case it finalizes all other concurrently
  running instances and outputs its own element from the batch.

  batched_tensor: The possibly transformed output of Batch. The size of the first
   dimension should remain unchanged by the transformations for the operation to
   work.
  batch_index: The matching batch_index obtained from Batch.
  id: The id scalar emitted by Batch.
  unbatched_tensor: The Tensor corresponding to this execution.
  timeout_micros: Maximum amount of time (in microseconds) to wait to receive the
   batched input tensor associated with a given invocation of the op.
  container: Container to control resource sharing.
  shared_name: Instances of Unbatch with the same container and shared_name are
   assumed to possibly belong to the same batch. If left empty, the op name will
   be used as the shared name.

  Args:
    batched_tensor: A `Tensor`.
    batch_index: A `Tensor` of type `int64`.
    id: A `Tensor` of type `int64`.
    timeout_micros: An `int`.
    container: An optional `string`. Defaults to `""`.
    shared_name: An optional `string`. Defaults to `""`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor`. Has the same type as `batched_tensor`.
  """
  _ctx = _context._context or _context.context()
  tld = _ctx._thread_local_data
  if tld.is_eager:
    try:
      _result = pywrap_tfe.TFE_Py_FastPathExecute(
        _ctx, "Unbatch", name, batched_tensor, batch_index, id,
        "timeout_micros", timeout_micros, "container", container,
        "shared_name", shared_name)
      return _result
    except _core._NotOkStatusException as e:
      _ops.raise_from_not_ok_status(e, name)
    except _core._FallbackException:
      pass
    try:
      return unbatch_eager_fallback(
          batched_tensor, batch_index, id, timeout_micros=timeout_micros,
          container=container, shared_name=shared_name, name=name, ctx=_ctx)
    except _core._SymbolicException:
      pass  # Add nodes to the TensorFlow graph.
  # Add nodes to the TensorFlow graph.
  timeout_micros = _execute.make_int(timeout_micros, "timeout_micros")
  if container is None:
    container = ""
  container = _execute.make_str(container, "container")
  if shared_name is None:
    shared_name = ""
  shared_name = _execute.make_str(shared_name, "shared_name")
  _, _, _op, _outputs = _op_def_library._apply_op_helper(
        "Unbatch", batched_tensor=batched_tensor, batch_index=batch_index,
                   id=id, timeout_micros=timeout_micros, container=container,
                   shared_name=shared_name, name=name)
  _result = _outputs[:]
  if _execute.must_record_gradient():
    _attrs = ("timeout_micros", _op._get_attr_int("timeout_micros"),
              "container", _op.get_attr("container"), "shared_name",
              _op.get_attr("shared_name"), "T", _op._get_attr_type("T"))
    _inputs_flat = _op.inputs
    _execute.record_gradient(
        "Unbatch", _inputs_flat, _attrs, _result)
  _result, = _result
  return _result

Unbatch = tf_export("raw_ops.Unbatch")(_ops.to_raw_op(unbatch))


def unbatch_eager_fallback(batched_tensor, batch_index, id, timeout_micros, container, shared_name, name, ctx):
  timeout_micros = _execute.make_int(timeout_micros, "timeout_micros")
  if container is None:
    container = ""
  container = _execute.make_str(container, "container")
  if shared_name is None:
    shared_name = ""
  shared_name = _execute.make_str(shared_name, "shared_name")
  _attr_T, (batched_tensor,) = _execute.args_to_matching_eager([batched_tensor], ctx, [])
  batch_index = _ops.convert_to_tensor(batch_index, _dtypes.int64)
  id = _ops.convert_to_tensor(id, _dtypes.int64)
  _inputs_flat = [batched_tensor, batch_index, id]
  _attrs = ("timeout_micros", timeout_micros, "container", container,
  "shared_name", shared_name, "T", _attr_T)
  _result = _execute.execute(b"Unbatch", 1, inputs=_inputs_flat, attrs=_attrs,
                             ctx=ctx, name=name)
  if _execute.must_record_gradient():
    _execute.record_gradient(
        "Unbatch", _inputs_flat, _attrs, _result)
  _result, = _result
  return _result


def unbatch_grad(original_input, batch_index, grad, id, container="", shared_name="", name=None):
  r"""Gradient of Unbatch.

  Acts like Batch but using the given batch_index index of batching things as they
  become available. This ensures that the gradients are propagated back in the
  same session which did the forward pass.

  original_input: The input to the Unbatch operation this is the gradient of.
  batch_index: The batch_index given to the Unbatch operation this is the gradient
  of.
  grad: The downstream gradient.
  id: The id scalar emitted by Batch.
  batched_grad: The return value, either an empty tensor or the batched gradient.
  container: Container to control resource sharing.
  shared_name: Instances of UnbatchGrad with the same container and shared_name
   are assumed to possibly belong to the same batch. If left empty, the op name
   will be used as the shared name.

  Args:
    original_input: A `Tensor`.
    batch_index: A `Tensor` of type `int64`.
    grad: A `Tensor`. Must have the same type as `original_input`.
    id: A `Tensor` of type `int64`.
    container: An optional `string`. Defaults to `""`.
    shared_name: An optional `string`. Defaults to `""`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor`. Has the same type as `original_input`.
  """
  _ctx = _context._context or _context.context()
  tld = _ctx._thread_local_data
  if tld.is_eager:
    try:
      _result = pywrap_tfe.TFE_Py_FastPathExecute(
        _ctx, "UnbatchGrad", name, original_input, batch_index, grad, id,
        "container", container, "shared_name", shared_name)
      return _result
    except _core._NotOkStatusException as e:
      _ops.raise_from_not_ok_status(e, name)
    except _core._FallbackException:
      pass
    try:
      return unbatch_grad_eager_fallback(
          original_input, batch_index, grad, id, container=container,
          shared_name=shared_name, name=name, ctx=_ctx)
    except _core._SymbolicException:
      pass  # Add nodes to the TensorFlow graph.
  # Add nodes to the TensorFlow graph.
  if container is None:
    container = ""
  container = _execute.make_str(container, "container")
  if shared_name is None:
    shared_name = ""
  shared_name = _execute.make_str(shared_name, "shared_name")
  _, _, _op, _outputs = _op_def_library._apply_op_helper(
        "UnbatchGrad", original_input=original_input, batch_index=batch_index,
                       grad=grad, id=id, container=container,
                       shared_name=shared_name, name=name)
  _result = _outputs[:]
  if _execute.must_record_gradient():
    _attrs = ("container", _op.get_attr("container"), "shared_name",
              _op.get_attr("shared_name"), "T", _op._get_attr_type("T"))
    _inputs_flat = _op.inputs
    _execute.record_gradient(
        "UnbatchGrad", _inputs_flat, _attrs, _result)
  _result, = _result
  return _result

UnbatchGrad = tf_export("raw_ops.UnbatchGrad")(_ops.to_raw_op(unbatch_grad))


def unbatch_grad_eager_fallback(original_input, batch_index, grad, id, container, shared_name, name, ctx):
  if container is None:
    container = ""
  container = _execute.make_str(container, "container")
  if shared_name is None:
    shared_name = ""
  shared_name = _execute.make_str(shared_name, "shared_name")
  _attr_T, _inputs_T = _execute.args_to_matching_eager([original_input, grad], ctx, [])
  (original_input, grad) = _inputs_T
  batch_index = _ops.convert_to_tensor(batch_index, _dtypes.int64)
  id = _ops.convert_to_tensor(id, _dtypes.int64)
  _inputs_flat = [original_input, batch_index, grad, id]
  _attrs = ("container", container, "shared_name", shared_name, "T", _attr_T)
  _result = _execute.execute(b"UnbatchGrad", 1, inputs=_inputs_flat,
                             attrs=_attrs, ctx=ctx, name=name)
  if _execute.must_record_gradient():
    _execute.record_gradient(
        "UnbatchGrad", _inputs_flat, _attrs, _result)
  _result, = _result
  return _result