def setUp(self):
    super(Base, self).setUp()

    self.x_size = 7
    self.channel_size = 3
    self.z_size = 4
    self.probs_size = 11

    tensor = math_ops.range(0, self.x_size * self.channel_size * self.z_size *
                            self.probs_size)
    tensor = array_ops.reshape(
        tensor, [self.x_size, self.channel_size, self.z_size, self.probs_size])
    a0 = ('x', range(self.x_size))
    a1 = ('channel', ['red', 'green', 'blue'])
    a2 = 'z'
    a3 = ('probs', np.linspace(0.0, 1.0, self.probs_size))

    self.tensor = tensor
    self.a0 = a0
    self.a1 = a1
    self.a2 = a2
    self.a2_resolved = ('z', self.z_size)
    self.a3 = a3
    self.original_lt = core.LabeledTensor(tensor, [a0, a1, a2, a3])

    self.x_probs_lt = core.slice_function(self.original_lt, {'z': 0})
    self.x_probs_lt = ops.select(self.x_probs_lt, {'channel': 'red'})
    self.channel_probs_lt = core.slice_function(self.original_lt,
                                                {'x': 3,
                                                 'z': 0})

  def test_slice_unlabeled(self):
    select_lt = core.slice_function(self.original_lt, {'z': slice(1, 3)})

    a2_sliced = 'z'
    golden_lt = core.LabeledTensor(self.tensor[:, :, 1:3, :],
                                   [self.a0, self.a1, a2_sliced, self.a3])

    self.assertLabeledTensorsEqual(select_lt, golden_lt)

  def test_slice(self):
    select_lt = core.slice_function(self.original_lt, {'channel': slice(0, 2)})

    a1_sliced = ('channel', ['red', 'green'])
    golden_lt = core.LabeledTensor(self.tensor[:, :2, :, :],
                                   [self.a0, a1_sliced, self.a2, self.a3])

    self.assertLabeledTensorsEqual(select_lt, golden_lt)

  def test_slices(self):
    select_lt = core.slice_function(self.original_lt, {'x': slice(1, 5),
                                                       'channel': slice(1,
                                                                        None)})

    a0_sliced = ('x', range(1, 5))
    a1_sliced = ('channel', ['green', 'blue'])
    golden_lt = core.LabeledTensor(self.tensor[1:5, 1:, :, :],
                                   [a0_sliced, a1_sliced, self.a2, self.a3])

    self.assertLabeledTensorsEqual(select_lt, golden_lt)

 def test(self):
   squeeze_lt = ops.squeeze(self.squeezable_lt, ['probs'])
   golden_lt = core.slice_function(self.squeezable_lt, {'probs': 0})
   self.assertLabeledTensorsEqual(squeeze_lt, golden_lt)

 def test_slice(self):
   map_lt = ops.map_fn(lambda t: core.slice_function(t, {'channel': 1}),
                       self.original_lt)
   slice_lt = core.slice_function(self.original_lt, {'channel': 1})
   self.assertLabeledTensorsEqual(map_lt, slice_lt)

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)

      else:
        # We're allowing anything NumPy treats as a scalar or 1D array.
        value = np.asarray(value)
        if value.ndim == 0:
          slices[axis_name] = axis.index(value.item())
        elif value.ndim == 1:
          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.tolist()]
          indexers[axis_name] = ops.convert_to_tensor(
              indexer, dtype=dtypes.int64)
        else:
          raise NotImplementedError(
              'select does not yet support selections with more than one '
              'dimension: %s on axis %r' % (value, axis_name))

    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)

 def test_slice_unknown_shape(self):
   lt = core.LabeledTensor(tf.placeholder(tf.float32, [None, 1]), ['x', 'y'])
   sliced_lt = core.slice_function(lt, {'y': 0})
   self.assertEqual(list(sliced_lt.axes.values()), [lt.axes['x']])

  def test_scalar(self):
    select_lt = core.slice_function(self.original_lt, {'channel': 1})
    golden_lt = core.LabeledTensor(self.tensor[:, 1, :, :], [self.a0, self.a2,
                                                             self.a3])

    self.assertLabeledTensorsEqual(select_lt, golden_lt)

 def test_name(self):
   select_lt = core.slice_function(self.original_lt, {'channel': 1})
   self.assertIn('lt_slice', select_lt.name)

 def test_name(self):
   foldl_lt = ops.foldl(core.add, self.original_lt,
                        core.slice_function(self.original_lt, {'x': 0}))
   self.assertIn('lt_foldl', foldl_lt.name)