# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Ops to manipulate lists of tensors."""

# pylint: disable=g-bad-name
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_list_ops
# go/tf-wildcard-import
# pylint: disable=wildcard-import
from tensorflow.python.ops.gen_list_ops import *
# pylint: enable=wildcard-import


ops.NotDifferentiable("TensorListConcatLists")
ops.NotDifferentiable("TensorListElementShape")
ops.NotDifferentiable("TensorListLength")
ops.NotDifferentiable("TensorListPushBackBatch")


def empty_tensor_list(element_shape,
                      element_dtype,
                      max_num_elements=None,
                      name=None):
  if max_num_elements is None:
    max_num_elements = -1

  return gen_list_ops.empty_tensor_list(
      element_shape=_build_element_shape(element_shape),
      element_dtype=element_dtype,
      max_num_elements=max_num_elements,
      name=name)


def tensor_list_reserve(element_shape, num_elements, element_dtype, name=None):
  return gen_list_ops.tensor_list_reserve(
      element_shape=_build_element_shape(element_shape),
      num_elements=num_elements,
      element_dtype=element_dtype,
      name=name)


def tensor_list_from_tensor(tensor, element_shape, name=None):
  return gen_list_ops.tensor_list_from_tensor(
      tensor=tensor,
      element_shape=_build_element_shape(element_shape),
      name=name)


def tensor_list_concat(input_handle, element_dtype, name=None):
  # Ignore the lengths output of TensorListConcat. It is only used during
  # gradient computation.
  return gen_list_ops.tensor_list_concat(
      input_handle=input_handle, element_dtype=element_dtype, name=name)[0]


def tensor_list_split(tensor, element_shape, lengths, name=None):
  return gen_list_ops.tensor_list_split(
      tensor=tensor,
      element_shape=_build_element_shape(element_shape),
      lengths=lengths,
      name=name)


@ops.RegisterGradient("TensorListPushBack")
def _PushBackGrad(op, dresult):
  return gen_list_ops.tensor_list_pop_back(
      dresult, element_dtype=op.get_attr("element_dtype"))


@ops.RegisterGradient("TensorListPopBack")
def _PopBackGrad(op, dlist, delement):
  if dlist is None:
    dlist = empty_tensor_list(
        element_dtype=delement.dtype,
        element_shape=gen_list_ops.tensor_list_element_shape(
            op.outputs[0], shape_type=dtypes.int32))
  return gen_list_ops.tensor_list_push_back(dlist, delement)


@ops.RegisterGradient("TensorListStack")
def _TensorListStackGrad(unused_op, dtensor):
  return tensor_list_from_tensor(dtensor, element_shape=dtensor.shape[1:])


@ops.RegisterGradient("TensorListConcat")
def _TensorListConcatGrad(op, dtensor, unused_dlengths):
  # TODO(srbs): We lose the element_shape information in tensor_list_concat.
  # Consider providing that as an output of TensorListConcat?
  if dtensor.shape.rank is None:
    element_shape = None
  else:
    element_shape = [None] + dtensor.shape.as_list()[1:]
  return tensor_list_split(
      dtensor,
      element_shape=_build_element_shape(element_shape),
      lengths=op.outputs[1])


@ops.RegisterGradient("TensorListSplit")
def _TensorListSplitGrad(op, dlist):
  return tensor_list_concat(dlist, element_dtype=op.inputs[0].dtype), None, None


@ops.RegisterGradient("TensorListFromTensor")
def _TensorListFromTensorGrad(op, dlist):
  """Gradient for TensorListFromTensor."""
  if op.inputs[0].shape.dims and op.inputs[0].shape.dims[0].value is not None:
    num_elements = op.inputs[0].shape.dims[0].value
  else:
    num_elements = None
  if dlist is None:
    dlist = empty_tensor_list(
        element_dtype=op.inputs[0].dtype,
        element_shape=gen_list_ops.tensor_list_element_shape(
            op.outputs[0], shape_type=dtypes.int32))
  tensor_grad = gen_list_ops.tensor_list_stack(
      dlist, element_dtype=op.inputs[0].dtype, num_elements=num_elements)
  shape_grad = None
  return tensor_grad, shape_grad


@ops.RegisterGradient("TensorListGetItem")
def _TensorListGetItemGrad(op, ditem):
  """Gradient for TensorListGetItem."""
  list_size = gen_list_ops.tensor_list_length(op.inputs[0])
  list_grad = gen_list_ops.tensor_list_set_item(
      gen_list_ops.tensor_list_reserve(
          gen_list_ops.tensor_list_element_shape(op.inputs[0],
                                                 shape_type=dtypes.int32),
          list_size, element_dtype=ditem.dtype),
      index=op.inputs[1],
      item=ditem)
  index_grad = None
  return list_grad, index_grad


@ops.RegisterGradient("TensorListSetItem")
def _TensorListSetItemGrad(op, dlist):
  _, index, item = op.inputs
  list_grad = gen_list_ops.tensor_list_set_item(
      dlist, index=index, item=array_ops.zeros_like(item))
  index_grad = None
  element_grad = gen_list_ops.tensor_list_get_item(
      dlist, index, element_dtype=item.dtype)
  return list_grad, index_grad, element_grad


@ops.RegisterGradient("TensorListGather")
def _TensorListGatherGrad(op, dtensor):
  _, indices = op.inputs
  return gen_list_ops.tensor_list_scatter(
      tensor=dtensor, indices=indices,
      element_shape=ops.convert_to_tensor(-1, dtype=dtypes.int32)), None


@ops.RegisterGradient("TensorListScatter")
def _TensorListScatterGrad(op, dlist):
  t, indices, _ = op.inputs
  return gen_list_ops.tensor_list_gather(
      dlist, indices, element_dtype=t.dtype), None


def _build_element_shape(shape):
  """Converts shape to a format understood by list_ops for element_shape.

  If `shape` is already a `Tensor` it is returned as-is. We do not perform a
  type check here.

  If shape is None or a TensorShape with unknown rank, -1 is returned.

  If shape is a scalar, an int32 tensor with empty list is returned. Note we
  do directly return an empty list since ops.convert_to_tensor would conver it
  to a float32 which is not a valid type for element_shape.

  If shape is a sequence of dims, None's in the list are replaced with -1. We
  do not check the dtype of the other dims.

  Args:
    shape: Could be None, Tensor, TensorShape or a list of dims (each dim could
      be a None, scalar or Tensor).

  Returns:
    A None-free shape that can be converted to a tensor.
  """
  if isinstance(shape, ops.Tensor):
    return shape
  if isinstance(shape, tensor_shape.TensorShape):
    # `TensorShape.as_list` requires rank to be known.
    shape = shape.as_list() if shape else None
  # Shape is unknown.
  if shape is None:
    return -1
  # Shape is a scalar.
  if not shape:
    return ops.convert_to_tensor(shape, dtype=dtypes.int32)
  # Shape is a sequence of dimensions. Convert None dims to -1.
  return [d if d is not None else -1 for d in shape]

[((33, 0, 33, 46), 'tensorflow.python.framework.ops.NotDifferentiable', 'ops.NotDifferentiable', ({(33, 22, 33, 45): '"""TensorListConcatLists"""'}, {}), "('TensorListConcatLists')", False, 'from tensorflow.python.framework import ops\n'), ((34, 0, 34, 47), 'tensorflow.python.framework.ops.NotDifferentiable', 'ops.NotDifferentiable', ({(34, 22, 34, 46): '"""TensorListElementShape"""'}, {}), "('TensorListElementShape')", False, 'from tensorflow.python.framework import ops\n'), ((35, 0, 35, 41), 'tensorflow.python.framework.ops.NotDifferentiable', 'ops.NotDifferentiable', ({(35, 22, 35, 40): '"""TensorListLength"""'}, {}), "('TensorListLength')", False, 'from tensorflow.python.framework import ops\n'), ((36, 0, 36, 48), 'tensorflow.python.framework.ops.NotDifferentiable', 'ops.NotDifferentiable', ({(36, 22, 36, 47): '"""TensorListPushBackBatch"""'}, {}), "('TensorListPushBackBatch')", False, 'from tensorflow.python.framework import ops\n'), ((83, 1, 83, 43), 'tensorflow.python.framework.ops.RegisterGradient', 'ops.RegisterGradient', ({(83, 22, 83, 42): '"""TensorListPushBack"""'}, {}), "('TensorListPushBack')", False, 'from tensorflow.python.framework import ops\n'), ((89, 1, 89, 42), 'tensorflow.python.framework.ops.RegisterGradient', 'ops.RegisterGradient', ({(89, 22, 89, 41): '"""TensorListPopBack"""'}, {}), "('TensorListPopBack')", False, 'from tensorflow.python.framework import ops\n'), ((99, 1, 99, 40), 'tensorflow.python.framework.ops.RegisterGradient', 'ops.RegisterGradient', ({(99, 22, 99, 39): '"""TensorListStack"""'}, {}), "('TensorListStack')", False, 'from tensorflow.python.framework import ops\n'), ((104, 1, 104, 41), 'tensorflow.python.framework.ops.RegisterGradient', 'ops.RegisterGradient', ({(104, 22, 104, 40): '"""TensorListConcat"""'}, {}), "('TensorListConcat')", False, 'from tensorflow.python.framework import ops\n'), ((118, 1, 118, 40), 'tensorflow.python.framework.ops.RegisterGradient', 'ops.RegisterGradient', ({(118, 22, 118, 39): '"""TensorListSplit"""'}, {}), "('TensorListSplit')", False, 'from tensorflow.python.framework import ops\n'), ((123, 1, 123, 45), 'tensorflow.python.framework.ops.RegisterGradient', 'ops.RegisterGradient', ({(123, 22, 123, 44): '"""TensorListFromTensor"""'}, {}), "('TensorListFromTensor')", False, 'from tensorflow.python.framework import ops\n'), ((141, 1, 141, 42), 'tensorflow.python.framework.ops.RegisterGradient', 'ops.RegisterGradient', ({(141, 22, 141, 41): '"""TensorListGetItem"""'}, {}), "('TensorListGetItem')", False, 'from tensorflow.python.framework import ops\n'), ((156, 1, 156, 42), 'tensorflow.python.framework.ops.RegisterGradient', 'ops.RegisterGradient', ({(156, 22, 156, 41): '"""TensorListSetItem"""'}, {}), "('TensorListSetItem')", False, 'from tensorflow.python.framework import ops\n'), ((167, 1, 167, 41), 'tensorflow.python.framework.ops.RegisterGradient', 'ops.RegisterGradient', ({(167, 22, 167, 40): '"""TensorListGather"""'}, {}), "('TensorListGather')", False, 'from tensorflow.python.framework import ops\n'), ((175, 1, 175, 42), 'tensorflow.python.framework.ops.RegisterGradient', 'ops.RegisterGradient', ({(175, 22, 175, 41): '"""TensorListScatter"""'}, {}), "('TensorListScatter')", False, 'from tensorflow.python.framework import ops\n'), ((96, 9, 96, 60), 'tensorflow.python.ops.gen_list_ops.tensor_list_push_back', 'gen_list_ops.tensor_list_push_back', ({(96, 44, 96, 49): 'dlist', (96, 51, 96, 59): 'delement'}, {}), '(dlist, delement)', False, 'from tensorflow.python.ops import gen_list_ops\n'), ((135, 16, 136, 73), 'tensorflow.python.ops.gen_list_ops.tensor_list_stack', 'gen_list_ops.tensor_list_stack', (), '', False, 'from tensorflow.python.ops import gen_list_ops\n'), ((144, 14, 144, 59), 'tensorflow.python.ops.gen_list_ops.tensor_list_length', 'gen_list_ops.tensor_list_length', ({(144, 46, 144, 58): 'op.inputs[0]'}, {}), '(op.inputs[0])', False, 'from tensorflow.python.ops import gen_list_ops\n'), ((162, 17, 163, 45), 'tensorflow.python.ops.gen_list_ops.tensor_list_get_item', 'gen_list_ops.tensor_list_get_item', (), '', False, 'from tensorflow.python.ops import gen_list_ops\n'), ((71, 9, 72, 72), 'tensorflow.python.ops.gen_list_ops.tensor_list_concat', 'gen_list_ops.tensor_list_concat', (), '', False, 'from tensorflow.python.ops import gen_list_ops\n'), ((178, 9, 179, 44), 'tensorflow.python.ops.gen_list_ops.tensor_list_gather', 'gen_list_ops.tensor_list_gather', (), '', False, 'from tensorflow.python.ops import gen_list_ops\n'), ((214, 11, 214, 59), 'tensorflow.python.framework.ops.convert_to_tensor', 'ops.convert_to_tensor', (), '', False, 'from tensorflow.python.framework import ops\n'), ((147, 10, 148, 73), 'tensorflow.python.ops.gen_list_ops.tensor_list_element_shape', 'gen_list_ops.tensor_list_element_shape', (), '', False, 'from tensorflow.python.ops import gen_list_ops\n'), ((160, 31, 160, 57), 'tensorflow.python.ops.array_ops.zeros_like', 'array_ops.zeros_like', ({(160, 52, 160, 56): 'item'}, {}), '(item)', False, 'from tensorflow.python.ops import array_ops\n'), ((94, 22, 95, 51), 'tensorflow.python.ops.gen_list_ops.tensor_list_element_shape', 'gen_list_ops.tensor_list_element_shape', (), '', False, 'from tensorflow.python.ops import gen_list_ops\n'), ((133, 22, 134, 51), 'tensorflow.python.ops.gen_list_ops.tensor_list_element_shape', 'gen_list_ops.tensor_list_element_shape', (), '', False, 'from tensorflow.python.ops import gen_list_ops\n'), ((172, 20, 172, 65), 'tensorflow.python.framework.ops.convert_to_tensor', 'ops.convert_to_tensor', (), '', False, 'from tensorflow.python.framework import ops\n')]