# Copyright 2016 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.
# ==============================================================================
"""Non-core ops for LabeledTensor."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import types

import numpy as np
from six import string_types

from tensorflow.contrib.labeled_tensor.python.ops import _typecheck as tc
from tensorflow.contrib.labeled_tensor.python.ops import core
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import functional_ops
from tensorflow.python.ops import map_fn as map_fn_lib
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import numerics
from tensorflow.python.ops import random_ops
from tensorflow.python.training import input  # pylint: disable=redefined-builtin


@tc.returns(core.LabeledTensor)
@tc.accepts(core.LabeledTensor, ops.Tensor, core.Axis,
            tc.Optional(string_types))
def _gather_1d_on_axis(labeled_tensor, indexer, axis, name=None):
  with ops.name_scope(name, 'lt_take', [labeled_tensor]) as scope:
    temp_axes = core.Axes([axis] + list(
        labeled_tensor.axes.remove(axis.name).values()))
    transposed = core.transpose(labeled_tensor, temp_axes.keys())
    indexed = core.LabeledTensor(
        array_ops.gather(transposed.tensor, indexer), temp_axes)
    return core.transpose(indexed, labeled_tensor.axes.keys(), name=scope)


@tc.returns(core.LabeledTensor)
@tc.accepts(core.LabeledTensorLike,
            tc.Mapping(string_types,
                       tc.Union(slice, collections.Hashable, list)),
            tc.Optional(string_types))
def select(labeled_tensor, selection, name=None):
  """Slice out a subset of the tensor.

  Args:
    labeled_tensor: The input tensor.
    selection: A dictionary mapping an axis name to a scalar, slice or list of
      values to select. Currently supports two types of selections:
        (a) Any number of scalar and/or slice selections.
        (b) Exactly one list selection, without any scalars or slices.
    name: Optional op name.

  Returns:
    The selection as a `LabeledTensor`.

  Raises:
    ValueError: If the tensor doesn't have an axis in the selection or if
      that axis lacks labels.
    KeyError: If any labels in a selection are not found in the original axis.
    NotImplementedError: If you attempt to combine a list selection with
      scalar selection or another list selection.
  """
  with ops.name_scope(name, 'lt_select', [labeled_tensor]) as scope:
    labeled_tensor = core.convert_to_labeled_tensor(labeled_tensor)

    slices = {}
    indexers = {}
    for axis_name, value in selection.items():
      if axis_name not in labeled_tensor.axes:
        raise ValueError(
            'The tensor does not have an axis named %s. Its axes are: %r' %
            (axis_name, labeled_tensor.axes.keys()))
      axis = labeled_tensor.axes[axis_name]
      if axis.labels is None:
        raise ValueError(
            'The axis named %s does not have labels. The axis is: %r' %
            (axis_name, axis))

      if isinstance(value, slice):
        # TODO(shoyer): consider deprecating using slices in favor of lists
        if value.start is None:
          start = None
        else:
          start = axis.index(value.start)

        if value.stop is None:
          stop = None
        else:
          # For now, follow the pandas convention of making labeled slices
          # inclusive of both bounds.
          stop = axis.index(value.stop) + 1

        if value.step is not None:
          raise NotImplementedError('slicing with a step is not yet supported')

        slices[axis_name] = slice(start, stop)

      # Needs to be after checking for slices, since slice objects claim to be
      # instances of collections.Hashable but hash() on them fails.
      elif isinstance(value, collections.Hashable):
        slices[axis_name] = axis.index(value)

      elif isinstance(value, list):
        if indexers:
          raise NotImplementedError(
              'select does not yet support more than one list selection at '
              'the same time')
        indexer = [axis.index(v) for v in value]
        indexers[axis_name] = ops.convert_to_tensor(indexer, dtype=dtypes.int64)

      else:
        # If type checking is working properly, this shouldn't be possible.
        raise TypeError('cannot handle arbitrary types')

    if indexers and slices:
      raise NotImplementedError(
          'select does not yet support combined scalar and list selection')

    # For now, handle array selection separately, because tf.gather_nd does
    # not support gradients yet. Later, using gather_nd will let us combine
    # these paths.
    if indexers:
      (axis_name, indexer), = indexers.items()
      axis = core.Axis(axis_name, selection[axis_name])
      return _gather_1d_on_axis(labeled_tensor, indexer, axis, name=scope)
    else:
      return core.slice_function(labeled_tensor, slices, name=scope)


@tc.returns(core.LabeledTensor)
@tc.accepts(
    tc.Collection(core.LabeledTensorLike), string_types,
    tc.Optional(string_types))
def concat(labeled_tensors, axis_name, name=None):
  """Concatenate tensors along a dimension.

  See tf.concat.

  Args:
    labeled_tensors: A list of input LabeledTensors.
    axis_name: The name of the axis along which to concatenate.
    name: Optional op name.

  Returns:
    The concatenated tensor.
    The coordinate labels for the concatenation dimension are also concatenated,
    if they are available for every tensor.

  Raises:
    ValueError: If fewer than one tensor inputs is provided, if the tensors
      have incompatible axes, or if `axis_name` isn't the name of an axis.
  """
  with ops.name_scope(name, 'lt_concat', labeled_tensors) as scope:
    labeled_tensors = [
        core.convert_to_labeled_tensor(lt) for lt in labeled_tensors
    ]

    if len(labeled_tensors) < 1:
      raise ValueError('concat expects at least 1 tensor, but received %s' %
                       labeled_tensors)

    # All tensors must have these axes.
    axes_0 = labeled_tensors[0].axes
    axis_names = list(axes_0.keys())

    if axis_name not in axis_names:
      raise ValueError('%s not in %s' % (axis_name, axis_names))

    shared_axes = axes_0.remove(axis_name)

    tensors = [labeled_tensors[0].tensor]
    concat_axis_list = [axes_0[axis_name]]
    for labeled_tensor in labeled_tensors[1:]:
      current_shared_axes = labeled_tensor.axes.remove(axis_name)
      if current_shared_axes != shared_axes:
        # TODO(shoyer): add more specific checks about what went wrong,
        # including raising AxisOrderError when appropriate
        raise ValueError('Mismatched shared axes: the first tensor '
                         'had axes %r but this tensor has axes %r.' %
                         (shared_axes, current_shared_axes))

      # Accumulate the axis labels, if they're available.
      concat_axis_list.append(labeled_tensor.axes[axis_name])
      tensors.append(labeled_tensor.tensor)

    concat_axis = core.concat_axes(concat_axis_list)
    concat_dimension = axis_names.index(axis_name)
    concat_tensor = array_ops.concat(tensors, concat_dimension, name=scope)
    values = list(axes_0.values())
    concat_axes = (values[:concat_dimension] + [concat_axis] +
                   values[concat_dimension + 1:])

    return core.LabeledTensor(concat_tensor, concat_axes)


# TODO(shoyer): rename pack/unpack to stack/unstack


@tc.returns(core.LabeledTensor)
@tc.accepts(
    tc.Collection(core.LabeledTensorLike),
    tc.Union(string_types, core.AxisLike), int, tc.Optional(string_types))
def pack(labeled_tensors, new_axis, axis_position=0, name=None):
  """Pack tensors along a new axis.

  See tf.pack.

  Args:
    labeled_tensors: The input tensors, which must have identical axes.
    new_axis: The name of the new axis, or a tuple containing the name
      and coordinate labels.
    axis_position: Optional integer position at which to insert the new axis.
    name: Optional op name.

  Returns:
    The packed tensors as a single LabeledTensor, with `new_axis` in the given
    `axis_position`.

  Raises:
    ValueError: If fewer than one input tensors is provided, or if the tensors
      don't have identical axes.
  """
  with ops.name_scope(name, 'lt_pack', labeled_tensors) as scope:
    labeled_tensors = [
        core.convert_to_labeled_tensor(lt) for lt in labeled_tensors
    ]

    if len(labeled_tensors) < 1:
      raise ValueError('pack expects at least 1 tensors, but received %s' %
                       labeled_tensors)

    axes_0 = labeled_tensors[0].axes
    for t in labeled_tensors:
      if t.axes != axes_0:
        raise ValueError('Non-identical axes. Expected %s but got %s' %
                         (axes_0, t.axes))

    pack_op = array_ops.stack(
        [t.tensor for t in labeled_tensors], axis=axis_position, name=scope)
    axes = list(axes_0.values())
    axes.insert(axis_position, new_axis)
    return core.LabeledTensor(pack_op, axes)


@tc.returns(tc.List(core.LabeledTensor))
@tc.accepts(core.LabeledTensorLike,
            tc.Optional(string_types), tc.Optional(string_types))
def unpack(labeled_tensor, axis_name=None, name=None):
  """Unpack the tensor.

  See tf.unpack.

  Args:
    labeled_tensor: The input tensor.
    axis_name: Optional name of axis to unpack. By default, the first axis is
      used.
    name: Optional op name.

  Returns:
    The list of unpacked LabeledTensors.

  Raises:
    ValueError: If `axis_name` is not an axis on the input.
  """
  with ops.name_scope(name, 'lt_unpack', [labeled_tensor]) as scope:
    labeled_tensor = core.convert_to_labeled_tensor(labeled_tensor)

    axis_names = list(labeled_tensor.axes.keys())
    if axis_name is None:
      axis_name = axis_names[0]

    if axis_name not in axis_names:
      raise ValueError('%s not in %s' % (axis_name, axis_names))
    axis = axis_names.index(axis_name)

    unpack_ops = array_ops.unstack(labeled_tensor.tensor, axis=axis, name=scope)
    axes = [a for i, a in enumerate(labeled_tensor.axes.values()) if i != axis]
    return [core.LabeledTensor(t, axes) for t in unpack_ops]


@tc.returns(core.LabeledTensor)
@tc.accepts(core.LabeledTensorLike,
            tc.Collection(string_types),
            tc.Collection(tc.Union(string_types, core.AxisLike)),
            tc.Optional(string_types))
def reshape(labeled_tensor, existing_axes, new_axes, name=None):
  """Reshape specific axes of a LabeledTensor.

  Non-indicated axes remain in their original locations.

  Args:
    labeled_tensor: The input tensor.
    existing_axes: List of axis names found on the input tensor. These must
      appear sequentially in the list of axis names on the input. In other
      words, they must be a valid slice of `list(labeled_tensor.axes.keys())`.
    new_axes: List of strings, tuples of (axis_name, axis_value) or Axis objects
      providing new axes with which to replace `existing_axes` in the reshaped
      result. At most one element of `new_axes` may be a string, indicating an
      axis with unknown size.
    name: Optional op name.

  Returns:
    The reshaped LabeledTensor.

  Raises:
    ValueError: If `existing_axes` are not all axes on the input, or if more
     than one of `new_axes` has unknown size.
    AxisOrderError: If `existing_axes` are not a slice of axis names on the
      input.
  """
  with ops.name_scope(name, 'lt_reshape', [labeled_tensor]) as scope:
    labeled_tensor = core.convert_to_labeled_tensor(labeled_tensor)

    original_axis_names = list(labeled_tensor.axes.keys())
    existing_axes = list(existing_axes)
    if not set(existing_axes) <= set(original_axis_names):
      raise ValueError('existing_axes %r are not contained in the set of axis '
                       'names %r on the input labeled tensor' %
                       (existing_axes, original_axis_names))

    start = original_axis_names.index(existing_axes[0])
    stop = original_axis_names.index(existing_axes[-1]) + 1

    if existing_axes != original_axis_names[start:stop]:
      # We could support existing_axes that aren't a slice by using transpose,
      # but that could lead to unpredictable performance consequences because
      # transposes are not free in TensorFlow. If we did transpose
      # automatically, the user might never realize that their data is being
      # produced with the wrong order. (The later will occur with some frequency
      # because of how broadcasting automatically choose axis order.)
      # So for now we've taken the strict approach.
      raise core.AxisOrderError(
          'existing_axes %r are not a slice of axis names %r on the input '
          'labeled tensor. Use `transpose` or `impose_axis_order` to reorder '
          'axes on the input explicitly.' %
          (existing_axes, original_axis_names))

    if sum(isinstance(axis, string_types) for axis in new_axes) > 1:
      raise ValueError(
          'at most one axis in new_axes can have unknown size. All other '
          'axes must have an indicated integer size or labels: %r' % new_axes)

    original_values = list(labeled_tensor.axes.values())
    axis_size = lambda axis: -1 if axis.size is None else axis.size
    shape = [axis_size(axis) for axis in original_values[:start]]
    for axis_ref in new_axes:
      if isinstance(axis_ref, string_types):
        shape.append(-1)
      else:
        axis = core.as_axis(axis_ref)
        shape.append(axis_size(axis))
    shape.extend(axis_size(axis) for axis in original_values[stop:])

    reshaped_tensor = array_ops.reshape(
        labeled_tensor.tensor, shape, name=scope)
    axes = original_values[:start] + list(new_axes) + original_values[stop:]
    return core.LabeledTensor(reshaped_tensor, axes)


@tc.returns(core.LabeledTensor)
@tc.accepts(core.LabeledTensorLike, string_types, string_types,
            tc.Optional(string_types))
def rename_axis(labeled_tensor, existing_name, new_name, name=None):
  """Rename an axis of LabeledTensor.

  Args:
    labeled_tensor: The input tensor.
    existing_name: Name for an existing axis on the input.
    new_name: Desired replacement name.
    name: Optional op name.

  Returns:
    LabeledTensor with renamed axis.

  Raises:
    ValueError: If `existing_name` is not an axis on the input.
  """
  with ops.name_scope(name, 'lt_rename_axis', [labeled_tensor]) as scope:
    if existing_name not in labeled_tensor.axes:
      raise ValueError('existing_name %r are not contained in the set of axis '
                       'names %r on the input labeled tensor' %
                       (existing_name, labeled_tensor.axes.keys()))
    new_axis = core.Axis(new_name, labeled_tensor.axes[existing_name].value)
    return reshape(labeled_tensor, [existing_name], [new_axis], name=scope)


@tc.returns(tc.List(core.LabeledTensor))
@tc.accepts(string_types, collections.Callable, int, bool,
            tc.Collection(core.LabeledTensorLike), bool,
            tc.Optional(string_types))
def _batch_helper(default_name,
                  batch_fn,
                  batch_size,
                  enqueue_many,
                  labeled_tensors,
                  allow_smaller_final_batch,
                  name=None):
  with ops.name_scope(name, default_name, labeled_tensors) as scope:
    labeled_tensors = [
        core.convert_to_labeled_tensor(lt) for lt in labeled_tensors
    ]

    batch_ops = batch_fn([t.tensor for t in labeled_tensors], scope)
    # TODO(shoyer): Remove this when they sanitize the TF API.
    if not isinstance(batch_ops, list):
      assert isinstance(batch_ops, ops.Tensor)
      batch_ops = [batch_ops]

    if allow_smaller_final_batch:
      batch_size = None

    @tc.returns(core.Axes)
    @tc.accepts(core.Axes)
    def output_axes(axes):
      if enqueue_many:
        if 'batch' not in axes or list(axes.keys()).index('batch') != 0:
          raise ValueError(
              'When enqueue_many is True, input tensors must have an axis '
              'called "batch" as their first dimension, '
              'but axes were %s' % axes)
        culled_axes = axes.remove('batch')
        return core.Axes([('batch', batch_size)] + list(culled_axes.values()))
      else:
        return core.Axes([('batch', batch_size)] + list(axes.values()))

    output_labeled_tensors = []
    for i, tensor in enumerate(batch_ops):
      axes = output_axes(labeled_tensors[i].axes)
      output_labeled_tensors.append(core.LabeledTensor(tensor, axes))

    return output_labeled_tensors


@tc.returns(tc.List(core.LabeledTensor))
@tc.accepts(
    tc.Collection(core.LabeledTensorLike), int, int, int, bool, bool,
    tc.Optional(string_types))
def batch(labeled_tensors,
          batch_size,
          num_threads=1,
          capacity=32,
          enqueue_many=False,
          allow_smaller_final_batch=False,
          name=None):
  """Rebatch a tensor.

  See tf.batch.

  Args:
    labeled_tensors: The input tensors.
    batch_size: The output batch size.
    num_threads: See tf.batch.
    capacity: See tf.batch.
    enqueue_many: If true, the input tensors must contain a 'batch' axis as
      their first axis.
      If false, the input tensors must not contain a 'batch' axis.
      See tf.batch.
    allow_smaller_final_batch: See tf.batch.
    name: Optional op name.

  Returns:
    The rebatched tensors.
    If enqueue_many is false, the output tensors will have a new 'batch' axis
      as their first axis.

  Raises:
    ValueError: If enqueue_many is True and the first axis of the tensors
      isn't "batch".
  """

  def fn(tensors, scope):
    return input.batch(
        tensors,
        batch_size=batch_size,
        num_threads=num_threads,
        capacity=capacity,
        enqueue_many=enqueue_many,
        allow_smaller_final_batch=allow_smaller_final_batch,
        name=scope)

  return _batch_helper('lt_batch', fn, batch_size, enqueue_many,
                       labeled_tensors, allow_smaller_final_batch, name)


@tc.returns(tc.List(core.LabeledTensor))
@tc.accepts(
    tc.Collection(core.LabeledTensorLike), int, int, int, bool, int,
    tc.Optional(int), bool, tc.Optional(string_types))
def shuffle_batch(labeled_tensors,
                  batch_size,
                  num_threads=1,
                  capacity=32,
                  enqueue_many=False,
                  min_after_dequeue=0,
                  seed=None,
                  allow_smaller_final_batch=False,
                  name=None):
  """Rebatch a tensor, with shuffling.

  See tf.batch.

  Args:
    labeled_tensors: The input tensors.
    batch_size: The output batch size.
    num_threads: See tf.batch.
    capacity: See tf.batch.
    enqueue_many: If true, the input tensors must contain a 'batch' axis as
      their first axis.
      If false, the input tensors must not contain a 'batch' axis.
      See tf.batch.
    min_after_dequeue: Minimum number of elements in the queue after a dequeue,
      used to ensure mixing.
    seed: Optional random seed.
    allow_smaller_final_batch: See tf.batch.
    name: Optional op name.

  Returns:
    The rebatched tensors.
    If enqueue_many is false, the output tensors will have a new 'batch' axis
      as their first axis.

  Raises:
    ValueError: If enqueue_many is True and the first axis of the tensors
      isn't "batch".
  """

  def fn(tensors, scope):
    return input.shuffle_batch(
        tensors,
        batch_size=batch_size,
        num_threads=num_threads,
        capacity=capacity,
        enqueue_many=enqueue_many,
        min_after_dequeue=min_after_dequeue,
        seed=seed,
        allow_smaller_final_batch=allow_smaller_final_batch,
        name=scope)

  return _batch_helper('lt_shuffle_batch', fn, batch_size, enqueue_many,
                       labeled_tensors, allow_smaller_final_batch, name)


@tc.returns(core.LabeledTensor)
@tc.accepts(core.LabeledTensorLike,
            tc.Mapping(string_types, int),
            tc.Optional(int), tc.Optional(string_types))
def random_crop(labeled_tensor, shape_map, seed=None, name=None):
  """Randomly crops a tensor to a given size.

  See tf.random_crop.

  Args:
    labeled_tensor: The input tensor.
    shape_map: A dictionary mapping axis names to the size of the random crop
      for that dimension.
    seed: An optional random seed.
    name: An optional op name.

  Returns:
    A tensor of the same rank as `labeled_tensor`, cropped randomly in the
    selected dimensions.

  Raises:
    ValueError: If the shape map contains an axis name not in the input tensor.
  """
  with ops.name_scope(name, 'lt_random_crop', [labeled_tensor]) as scope:
    labeled_tensor = core.convert_to_labeled_tensor(labeled_tensor)

    for axis_name in shape_map:
      if axis_name not in labeled_tensor.axes:
        raise ValueError('Selection axis %s not in axes %s' %
                         (axis_name, labeled_tensor.axes))

    shape = []
    axes = []
    for axis in labeled_tensor.axes.values():
      if axis.name in shape_map:
        size = shape_map[axis.name]
        shape.append(size)
        # We lose labels for the axes we crop, leaving just the size.
        axes.append((axis.name, size))
      else:
        shape.append(len(axis))
        axes.append(axis)

    crop_op = random_ops.random_crop(
        labeled_tensor.tensor, shape, seed=seed, name=scope)

    return core.LabeledTensor(crop_op, axes)


# TODO(shoyer): Allow the user to select the axis over which to map.
@tc.returns(core.LabeledTensor)
@tc.accepts(collections.Callable, core.LabeledTensorLike,
            tc.Optional(string_types))
def map_fn(fn, labeled_tensor, name=None):
  """Map on the list of tensors unpacked from labeled_tensor.

  See tf.map_fn.

  Args:
    fn: The function to apply to each unpacked LabeledTensor.
      It should have type LabeledTensor -> LabeledTensor.
    labeled_tensor: The input tensor.
    name: Optional op name.

  Returns:
    A tensor that packs the results of applying fn to the list of tensors
    unpacked from labeled_tensor.
  """
  with ops.name_scope(name, 'lt_map_fn', [labeled_tensor]) as scope:
    labeled_tensor = core.convert_to_labeled_tensor(labeled_tensor)

    unpack_lts = unpack(labeled_tensor)

    # TODO(ericmc): Fix this upstream.
    if labeled_tensor.dtype == dtypes.string:
      # We must construct the full graph here, because map_fn_lib.map_fn
      # doesn't work for string-valued tensors.
      # Constructing the full graph may be slow.
      map_lts = [fn(t) for t in unpack_lts]
      return pack(map_lts, list(labeled_tensor.axes.values())[0], name=scope)
    else:
      # Figure out what the axis labels should be, but use tf.map_fn to
      # construct the graph because it's efficient.
      # It may be slow to construct the full graph, so we infer the labels from
      # the first element.
      # TODO(ericmc): This builds a subgraph which then gets thrown away.
      # Find a more elegant solution.
      first_map_lt = fn(unpack_lts[0])
      final_axes = list(labeled_tensor.axes.values())[:1] + list(
          first_map_lt.axes.values())

      @tc.returns(ops.Tensor)
      @tc.accepts(ops.Tensor)
      def tf_fn(tensor):
        original_axes = list(labeled_tensor.axes.values())[1:]
        tensor_lt = core.LabeledTensor(tensor, original_axes)
        return fn(tensor_lt).tensor

      map_op = map_fn_lib.map_fn(
          tf_fn, labeled_tensor.tensor, dtype=first_map_lt.dtype)
      map_lt = core.LabeledTensor(map_op, final_axes)

      return core.identity(map_lt, name=scope)


@tc.returns(core.LabeledTensor)
@tc.accepts(collections.Callable, core.LabeledTensorLike,
            core.LabeledTensorLike, tc.Optional(string_types))
def foldl(fn, labeled_tensor, initial_value, name=None):
  """Left fold on the list of tensors unpacked from labeled_tensor.

  See tf.foldl.

  Args:
    fn: The function to apply to each unpacked LabeledTensor.
      It should have type (LabeledTensor, LabeledTensor) -> LabeledTensor.
      Its arguments are (accumulated_value, next_value).
    labeled_tensor: The input tensor.
    initial_value: The initial value of the accumulator.
    name: Optional op name.

  Returns:
    The accumulated value.
  """
  with ops.name_scope(name, 'lt_foldl',
                      [labeled_tensor, initial_value]) as scope:
    labeled_tensor = core.convert_to_labeled_tensor(labeled_tensor)
    initial_value = core.convert_to_labeled_tensor(initial_value)

    @tc.returns(ops.Tensor)
    @tc.accepts(ops.Tensor, ops.Tensor)
    def tf_fn(accumulator, next_element):
      accumulator_lt = core.LabeledTensor(accumulator, initial_value.axes)
      next_element_lt = core.LabeledTensor(
          next_element, list(labeled_tensor.axes.values())[1:])
      return fn(accumulator_lt, next_element_lt).tensor

    foldl_op = functional_ops.foldl(
        tf_fn, labeled_tensor.tensor, initializer=initial_value.tensor)
    foldl_lt = core.LabeledTensor(foldl_op, initial_value.axes)

    return core.identity(foldl_lt, name=scope)


@tc.returns(core.LabeledTensor)
@tc.accepts(core.LabeledTensorLike,
            tc.Optional(tc.Collection(string_types)), tc.Optional(string_types))
def squeeze(labeled_tensor, axis_names=None, name=None):
  """Remove size-1 dimensions.

  See tf.squeeze.

  Args:
    labeled_tensor: The input tensor.
    axis_names: The names of the dimensions to remove, or None to remove
      all size-1 dimensions.
    name: Optional op name.

  Returns:
    A tensor with the specified dimensions removed.

  Raises:
    ValueError: If the named axes are not in the tensor, or if they are
      not size-1.
  """
  with ops.name_scope(name, 'lt_squeeze', [labeled_tensor]) as scope:
    labeled_tensor = core.convert_to_labeled_tensor(labeled_tensor)

    if axis_names is None:
      axis_names = [a.name for a in labeled_tensor.axes.values() if len(a) == 1]

    for axis_name in axis_names:
      if axis_name not in labeled_tensor.axes:
        raise ValueError('axis %s is not in tensor axes %s' %
                         (axis_name, labeled_tensor.axes))
      elif len(labeled_tensor.axes[axis_name]) != 1:
        raise ValueError(
            'cannot squeeze axis with size greater than 1: (%s, %s)' %
            (axis_name, labeled_tensor.axes[axis_name]))

    squeeze_dimensions = []
    axes = []
    for i, axis in enumerate(labeled_tensor.axes.values()):
      if axis.name in axis_names:
        squeeze_dimensions.append(i)
      else:
        axes.append(axis)

    if squeeze_dimensions:
      squeeze_op = array_ops.squeeze(
          labeled_tensor.tensor, squeeze_dimensions, name=scope)
    else:
      squeeze_op = array_ops.identity(labeled_tensor.tensor, name=scope)

    return core.LabeledTensor(squeeze_op, axes)


# pylint: disable=invalid-name
ReduceAxis = tc.Union(string_types,
                      tc.Tuple(string_types, collections.Hashable))
ReduceAxes = tc.Optional(tc.Union(ReduceAxis, tc.Collection(ReduceAxis)))
# pylint: enable=invalid-name


@tc.returns(core.LabeledTensor)
@tc.accepts(core.LabeledTensorLike, core.LabeledTensorLike,
            tc.Optional(string_types))
def matmul(a, b, name=None):
  """Matrix multiply two tensors with rank 1 or 2.

  If both tensors have rank 2, a matrix-matrix product is performed.
  If one tensor has rank 1 and the other has rank 2, then a matrix-vector
  product is performed.
  If both tensors have rank 1, then a vector dot-product is performed.
  (This behavior matches that of `numpy.dot`.)

  Both tensors must share exactly one dimension in common, which is the
  dimension the operation is summed along. The inputs will be automatically