def _diagonal_hessian(self, grads, predictions):
    """Prepares hessians for diagonal-hessian multiclass mode."""
    diag_hessian_list = []

    gradients_list = array_ops.unstack(
        grads, num=self._learner_config.num_classes, axis=1)

    for row, row_grads in enumerate(gradients_list):
      # If current row is i, K is number of classes,each row returns a tensor of
      # size batch_size x K representing for each example dx_i dx_1, dx_1 dx_2
      # etc dx_i dx_K
      hessian_row = gradients_impl.gradients(
          row_grads,
          predictions,
          name="Hessian_%d" % row,
          colocate_gradients_with_ops=False,
          gate_gradients=0,
          aggregation_method=None)

      # hessian_row is of dimension 1, batch_size, K, => trim first dimension
      # to get batch_size x K
      hessian_row = array_ops.squeeze(array_ops.unstack(hessian_row), [0])

      # Get dx_i^2 for the whole batch.
      elem = array_ops.transpose(hessian_row)[row]
      diag_hessian_list.append(elem)

    return diag_hessian_list

  def __call__(self,
               inputs,
               initial_state=None,
               dtype=None,
               sequence_length=None,
               scope=None):
    is_list = isinstance(inputs, list)
    if self._use_dynamic_rnn:
      if is_list:
        inputs = array_ops.stack(inputs)
      outputs, state = rnn.dynamic_rnn(
          self._cell,
          inputs,
          sequence_length=sequence_length,
          initial_state=initial_state,
          dtype=dtype,
          time_major=True,
          scope=scope)
      if is_list:
        # Convert outputs back to list
        outputs = array_ops.unstack(outputs)
    else:  # non-dynamic rnn
      if not is_list:
        inputs = array_ops.unstack(inputs)
      outputs, state = contrib_rnn.static_rnn(self._cell,
                                              inputs,
                                              initial_state=initial_state,
                                              dtype=dtype,
                                              sequence_length=sequence_length,
                                              scope=scope)
      if not is_list:
        # Convert outputs back to tensor
        outputs = array_ops.stack(outputs)

    return outputs, state

 def _add_comparison_summary(gan_model, reconstructions):
   image_list = (array_ops.unstack(gan_model.generator_inputs[:1]) +
                 array_ops.unstack(gan_model.generated_data[:1]) +
                 array_ops.unstack(reconstructions[:1]))
   summary.image(
       'image_comparison', eval_utils.image_reshaper(
           image_list, num_cols=len(image_list)), max_outputs=1)

  def _full_hessian(self, grads, predictions):
    """Prepares hessians for full-hessian multiclass strategy."""
    # Because of
    # https://github.com/tensorflow/tensorflow/issues/675, we can't just
    # compute the full hessian with a single call to gradients, but instead
    # must compute it row-by-row.
    gradients_list = array_ops.unstack(
        grads, num=self._learner_config.num_classes, axis=1)
    hessian_rows = []

    for row in range(self._learner_config.num_classes):
      # If current row is i, K is number of classes,each row returns a tensor of
      # size batch_size x K representing for each example dx_i dx_1, dx_i dx_2
      # etc dx_i dx_K
      hessian_row = gradients_impl.gradients(
          gradients_list[row],
          predictions,
          name="Hessian_%d" % row,
          colocate_gradients_with_ops=False,
          gate_gradients=0,
          aggregation_method=None)

      # hessian_row is of dimension 1, batch_size, K, => trim first dimension
      # to get batch_size x K
      hessian_row = array_ops.squeeze(array_ops.unstack(hessian_row), [0])
      hessian_rows.append(hessian_row)
    return hessian_rows

def seq2seq_inputs(x, y, input_length, output_length, sentinel=None, name=None):
  """Processes inputs for Sequence to Sequence models.

  Args:
    x: Input Tensor [batch_size, input_length, embed_dim].
    y: Output Tensor [batch_size, output_length, embed_dim].
    input_length: length of input x.
    output_length: length of output y.
    sentinel: optional first input to decoder and final output expected.
      If sentinel is not provided, zeros are used. Due to fact that y is not
      available in sampling time, shape of sentinel will be inferred from x.
    name: Operation name.

  Returns:
    Encoder input from x, and decoder inputs and outputs from y.
  """
  with ops.name_scope(name, "seq2seq_inputs", [x, y]):
    in_x = array_ops.unstack(x, axis=1)
    y = array_ops.unstack(y, axis=1)
    if not sentinel:
      # Set to zeros of shape of y[0], using x for batch size.
      sentinel_shape = array_ops.stack(
          [array_ops.shape(x)[0], y[0].get_shape()[1]])
      sentinel = array_ops.zeros(sentinel_shape)
      sentinel.set_shape(y[0].get_shape())
    in_y = [sentinel] + y
    out_y = y + [sentinel]
    return in_x, in_y, out_y

def _prepare_inputs_for_rnn(sequence_features, context_features, num_unroll):
  """Prepares features batched by the SQSS for input to a state-saving RNN.

  Args:
    sequence_features: A dict of sequence feature name to `Tensor`, with
      tensors of shape `[batch_size, num_unroll, ...]` and type float32.
    context_features: A dict of context feature name to `Tensor`, with
      tensors of shape `[batch_size, 1, ...]` and type float32.
    num_unroll: Python integer, how many time steps to unroll at a time.
      The input sequences of length `k` are then split into `k / num_unroll`
      many segments.

  Returns:
    features_by_time: A list of length `num_unroll` with `Tensor` entries of
      shape `[batch_size, len(sequence_features) + len(context_features)]` of
      type float32. Features are stored in lexicographic order by their
      corresponding feature dict keys, first in the `sequence_features` and
      then in the `context_features` dicts. Context features are copied into
      each time step.
  """

  def _tile(feature):
    return array_ops.squeeze(
        array_ops.tile(array_ops.expand_dims(feature, 1), [1, num_unroll, 1]),
        axis=2)

  sequence_features = [sequence_features[k] for k in sorted(sequence_features)]
  if not context_features:
    return array_ops.unstack(array_ops.stack(sequence_features, 2), axis=1)
  context_features = [
      _tile(context_features[k]) for k in sorted(context_features)
  ]
  return array_ops.unstack(
      array_ops.stack(sequence_features + context_features, 2), axis=1)

  def testAggregate(self):
    a = array_ops.constant([3., 4.])
    b = array_ops.constant([5., 6.])
    hint = op_hint.OpHint("agg")
    a0, a1 = array_ops.unstack(a)
    b0, b1 = array_ops.unstack(b)

    a0 = hint.add_input(a0, tag="c", aggregate=op_hint.OpHint.AGGREGATE_STACK)
    b0 = hint.add_input(b0, tag="n", aggregate=op_hint.OpHint.AGGREGATE_STACK)
    a1 = hint.add_input(a1, tag="c", aggregate=op_hint.OpHint.AGGREGATE_STACK)
    b1 = hint.add_input(b1, tag="n", aggregate=op_hint.OpHint.AGGREGATE_STACK)

    c0 = math_ops.add(a0, b0, name="addleft")
    c1 = math_ops.add(a1, b1, name="addright")
    c0 = hint.add_output(
        c0, tag="out", aggregate=op_hint.OpHint.AGGREGATE_STACK)
    c1 = hint.add_output(
        c1, tag="out", aggregate=op_hint.OpHint.AGGREGATE_STACK)

    curr = array_ops.stack([c0, c1])
    output = array_ops.identity(curr, name="FINAL_OUTPUT")
    with self.cached_session() as sess:
      stubbed_graphdef = op_hint.convert_op_hints_to_stubs(
          graph_def=sess.graph_def)
      self.assertEqual(
          self._getGraphOpTypes(
              stubbed_graphdef,
              output_nodes=[op_hint._tensor_name_base(output.name)]),
          set(["agg", "Const", "Identity"]))

 def testCannotInferNumFromNoneShape(self):
   x = array_ops.placeholder(np.float32, shape=(None,))
   with self.assertRaisesRegexp(ValueError,
                                r'Cannot infer num from shape \(\?,\)'):
     array_ops.unpack(x)
   with self.assertRaisesRegexp(ValueError,
                                r'Cannot infer num from shape \(\?,\)'):
     array_ops.unstack(x)

 def testCannotInferNumFromUnknownShape(self):
   x = array_ops.placeholder(np.float32)
   with self.assertRaisesRegexp(ValueError,
                                r'Cannot infer num from shape <unknown>'):
     array_ops.unpack(x)
   with self.assertRaisesRegexp(ValueError,
                                r'Cannot infer num from shape <unknown>'):
     array_ops.unstack(x)

def add_image_comparison_summaries(gan_model, num_comparisons=2,
                                   display_diffs=False):
  """Adds image summaries to compare triplets of images.

  The first image is the generator input, the second is the generator output,
  and the third is the real data. This style of comparison is useful for
  image translation problems, where the generator input is a corrupted image,
  the generator output is the reconstruction, and the real data is the target.

  Args:
    gan_model: A GANModel tuple.
    num_comparisons: The number of image triplets to display.
    display_diffs: Also display the difference between generated and target.

  Raises:
    ValueError: If real data, generated data, and generator inputs aren't
      images.
    ValueError: If the generator input, real, and generated data aren't all the
      same size.
  """
  if isinstance(gan_model, namedtuples.CycleGANModel):
    saved_params = locals()
    saved_params.pop('gan_model', None)
    with ops.name_scope('cyclegan_x2y_image_comparison_summaries'):
      add_image_comparison_summaries(gan_model.model_x2y, **saved_params)
    with ops.name_scope('cyclegan_y2x_image_comparison_summaries'):
      add_image_comparison_summaries(gan_model.model_y2x, **saved_params)
    return

  _assert_is_image(gan_model.generator_inputs)
  _assert_is_image(gan_model.generated_data)
  _assert_is_image(gan_model.real_data)

  gan_model.generated_data.shape.assert_is_compatible_with(
      gan_model.generator_inputs.shape)
  gan_model.real_data.shape.assert_is_compatible_with(
      gan_model.generated_data.shape)

  image_list = []
  image_list.extend(
      array_ops.unstack(gan_model.generator_inputs[:num_comparisons]))
  image_list.extend(
      array_ops.unstack(gan_model.generated_data[:num_comparisons]))
  image_list.extend(array_ops.unstack(gan_model.real_data[:num_comparisons]))
  if display_diffs:
    generated_list = array_ops.unstack(
        gan_model.generated_data[:num_comparisons])
    real_list = array_ops.unstack(gan_model.real_data[:num_comparisons])
    diffs = [
        math_ops.abs(math_ops.to_float(generated) - math_ops.to_float(real)) for
        generated, real in zip(generated_list, real_list)]
    image_list.extend(diffs)

  # Reshape image and display.
  summary.image(
      'image_comparison',
      eval_utils.image_reshaper(image_list, num_cols=num_comparisons),
      max_outputs=1)

def _prepare_inputs_for_rnn(sequence_features, context_features,
                            sequence_feature_columns, num_unroll):
  """Prepares features batched by the SQSS for input to a state-saving RNN.

  Args:
    sequence_features: A dict of sequence feature name to `Tensor` or
      `SparseTensor`, with `Tensor`s of shape `[batch_size, num_unroll, ...]`
      or `SparseTensors` of dense shape `[batch_size, num_unroll, d]`.
    context_features: A dict of context feature name to `Tensor`, with
      tensors of shape `[batch_size, 1, ...]` and type float32.
    sequence_feature_columns: An iterable containing all the feature columns
      describing sequence features. All items in the set should be instances
      of classes derived from `FeatureColumn`.
    num_unroll: Python integer, how many time steps to unroll at a time.
      The input sequences of length `k` are then split into `k / num_unroll`
      many segments.

  Returns:
    features_by_time: A list of length `num_unroll` with `Tensor` entries of
      shape `[batch_size, sum(sequence_features dimensions) +
      sum(context_features dimensions)]` of type float32.
      Context features are copied into each time step.
  """

  def _tile(feature):
    return array_ops.squeeze(
        array_ops.tile(array_ops.expand_dims(feature, 1), [1, num_unroll, 1]),
        axis=2)
  for feature in sequence_features.values():
    if isinstance(feature, sparse_tensor.SparseTensor):
      # Explicitly set dense_shape's shape to 3 ([batch_size, num_unroll, d])
      # since it can't be statically inferred.
      feature.dense_shape.set_shape([3])
  sequence_features = layers.sequence_input_from_feature_columns(
      columns_to_tensors=sequence_features,
      feature_columns=sequence_feature_columns,
      weight_collections=None,
      scope=None)
  # Explicitly set shape along dimension 1 to num_unroll for the unstack op.
  sequence_features.set_shape([None, num_unroll, None])

  if not context_features:
    return array_ops.unstack(sequence_features, axis=1)
  # TODO(jtbates): Call layers.input_from_feature_columns for context features.
  context_features = [
      _tile(context_features[k]) for k in sorted(context_features)
  ]
  return array_ops.unstack(
      array_ops.concat(
          [sequence_features, array_ops.stack(context_features, 2)], axis=2),
      axis=1)

 def _fn(x):
   """MADE parameterized via `masked_autoregressive_default_template`."""
   # TODO(b/67594795): Better support of dynamic shape.
   input_depth = x.shape.with_rank_at_least(1)[-1].value
   if input_depth is None:
     raise NotImplementedError(
         "Rightmost dimension must be known prior to graph execution.")
   input_shape = (np.int32(x.shape.as_list()) if x.shape.is_fully_defined()
                  else array_ops.shape(x))
   for i, units in enumerate(hidden_layers):
     x = masked_dense(
         inputs=x,
         units=units,
         num_blocks=input_depth,
         exclusive=True if i == 0 else False,
         activation=activation,
         *args,
         **kwargs)
   x = masked_dense(
       inputs=x,
       units=(1 if shift_only else 2) * input_depth,
       num_blocks=input_depth,
       activation=None,
       *args,
       **kwargs)
   if shift_only:
     x = array_ops.reshape(x, shape=input_shape)
     return x, None
   x = array_ops.reshape(
       x, shape=array_ops.concat([input_shape, [2]], axis=0))
   shift, log_scale = array_ops.unstack(x, num=2, axis=-1)
   which_clip = (math_ops.clip_by_value if log_scale_clip_gradient
                 else _clip_by_value_preserve_grad)
   log_scale = which_clip(log_scale, log_scale_min_clip, log_scale_max_clip)
   return shift, log_scale

def sequence_softmax(inputs, noutput, scope=None, name=None, linear_name=None):
  """Run a softmax layer over all the time steps of an input sequence.

  Args:
    inputs: (length, batch_size, depth) tensor
    noutput: output depth
    scope: optional scope name
    name: optional name for output tensor
    linear_name: name for linear (pre-softmax) output

  Returns:
    A tensor of size (length, batch_size, noutput).

  """
  length, _, ninputs = _shape(inputs)
  inputs_u = array_ops.unstack(inputs)
  output_u = []
  with variable_scope.variable_scope(scope, "SequenceSoftmax", [inputs]):
    initial_w = random_ops.truncated_normal([0 + ninputs, noutput], stddev=0.1)
    initial_b = constant_op.constant(0.1, shape=[noutput])
    w = variables.model_variable("weights", initializer=initial_w)
    b = variables.model_variable("biases", initializer=initial_b)
    for i in xrange(length):
      with variable_scope.variable_scope(scope, "SequenceSoftmaxStep",
                                         [inputs_u[i]]):
        # TODO(tmb) consider using slim.fully_connected(...,
        # activation_fn=tf.nn.softmax)
        linear = nn_ops.xw_plus_b(inputs_u[i], w, b, name=linear_name)
        output = nn_ops.softmax(linear)
        output_u += [output]
    outputs = array_ops.stack(output_u, name=name)
  return outputs

  def testBatch(self):
    # Build an arbitrary RGB image
    np.random.seed(7)
    batch_size = 5
    shape = (batch_size, 2, 7, 3)

    for nptype in self.float_types:
      inp = GenerateNumpyRandomRGB(shape).astype(nptype)

      # Convert to HSV and back, as a batch and individually
      with self.test_session() as sess:
        batch0 = array_ops.placeholder(nptype, shape=shape)
        with self.test_scope():
          batch1 = image_ops.rgb_to_hsv(batch0)
          batch2 = image_ops.hsv_to_rgb(batch1)
        split0 = array_ops.unstack(batch0)
        with self.test_scope():
          split1 = list(map(image_ops.rgb_to_hsv, split0))
          split2 = list(map(image_ops.hsv_to_rgb, split1))
        join1 = array_ops.stack(split1)
        join2 = array_ops.stack(split2)
        batch1, batch2, join1, join2 = sess.run([batch1, batch2, join1, join2],
                                                {batch0: inp})

      # Verify that processing batch elements together is the same as separate
      self.assertAllClose(batch1, join1)
      self.assertAllClose(batch2, join2)
      self.assertAllCloseAccordingToType(
          batch2, inp, bfloat16_atol=0.03, half_rtol=0.02)

  def test(self):
    unpack_lt = ops.unpack(self.original_lt)[0]
    golden_lt = core.LabeledTensor(
        array_ops.unstack(self.original_lt.tensor)[0],
        [self.a1, self.a2, self.a3])

    self.assertLabeledTensorsEqual(unpack_lt, golden_lt)

  def test_axis(self):
    unpack_lt = ops.unpack(self.original_lt, axis_name='z')[0]
    golden_lt = core.LabeledTensor(
        array_ops.unstack(
            self.original_lt.tensor, axis=2)[0], [self.a0, self.a1, self.a3])

    self.assertLabeledTensorsEqual(unpack_lt, golden_lt)

 def grow_tree_from_accumulated_summaries_fn():
   """Updates the tree with the best layer from accumulated summaries."""
   # Take out the accumulated summaries from the accumulator and grow.
   stats_summary_list = array_ops.unstack(
       summary_accumulator.take_grad(1), axis=0)
   grow_op = grow_tree_from_stats_summaries(stats_summary_list)
   return grow_op

def ndlstm_base_unrolled(inputs, noutput, scope=None, reverse=False):
  """Run an LSTM, either forward or backward.

  This is a 1D LSTM implementation using unrolling and the TensorFlow
  LSTM op.

  Args:
    inputs: input sequence (length, batch_size, ninput)
    noutput: depth of output
    scope: optional scope name
    reverse: run LSTM in reverse

  Returns:
    Output sequence (length, batch_size, noutput)

  """
  with variable_scope.variable_scope(scope, "SeqLstmUnrolled", [inputs]):
    length, batch_size, _ = _shape(inputs)
    lstm_cell = core_rnn_cell_impl.BasicLSTMCell(noutput, state_is_tuple=False)
    state = array_ops.zeros([batch_size, lstm_cell.state_size])
    output_u = []
    inputs_u = array_ops.unstack(inputs)
    if reverse:
      inputs_u = list(reversed(inputs_u))
    for i in xrange(length):
      if i > 0:
        variable_scope.get_variable_scope().reuse_variables()
      output, state = lstm_cell(inputs_u[i], state)
      output_u += [output]
    if reverse:
      output_u = list(reversed(output_u))
    outputs = array_ops.stack(output_u)
    return outputs

def _get_permutations(num_results, dims, seed=None):
  """Uniform iid sample from the space of permutations.

  Draws a sample of size `num_results` from the group of permutations of degrees
  specified by the `dims` tensor. These are packed together into one tensor
  such that each row is one sample from each of the dimensions in `dims`. For
  example, if dims = [2,3] and num_results = 2, the result is a tensor of shape
  [2, 2 + 3] and the first row of the result might look like:
  [1, 0, 2, 0, 1]. The first two elements are a permutation over 2 elements
  while the next three are a permutation over 3 elements.

  Args:
    num_results: A positive scalar `Tensor` of integral type. The number of
      draws from the discrete uniform distribution over the permutation groups.
    dims: A 1D `Tensor` of the same dtype as `num_results`. The degree of the
      permutation groups from which to sample.
    seed: (Optional) Python integer to seed the random number generator.

  Returns:
    permutations: A `Tensor` of shape `[num_results, sum(dims)]` and the same
    dtype as `dims`.
  """
  sample_range = math_ops.range(num_results)
  def generate_one(d):
    fn = lambda _: random_ops.random_shuffle(math_ops.range(d), seed=seed)
    return functional_ops.map_fn(fn, sample_range)
  return array_ops.concat([generate_one(d) for d in array_ops.unstack(dims)],
                          axis=-1)

  def transition_power_noise_accumulator(self, num_steps):
    """Computes power sums in closed form."""
    def _pack_and_reshape(*values):
      return array_ops.reshape(
          array_ops.stack(axis=1, values=values),
          array_ops.concat(values=[array_ops.shape(num_steps), [2, 2]], axis=0))

    num_steps = math_ops.cast(num_steps, self.dtype)
    noise_transitions = num_steps - 1
    noise_transform = ops.convert_to_tensor(self.get_noise_transform(),
                                            self.dtype)
    noise_covariance_transformed = math_ops.matmul(
        math_ops.matmul(noise_transform,
                        self.state_transition_noise_covariance),
        noise_transform,
        adjoint_b=True)
    # Un-packing the transformed noise as:
    # [[a b]
    #  [c d]]
    a, b, c, d = array_ops.unstack(
        array_ops.reshape(noise_covariance_transformed, [-1, 4]), axis=1)
    sum_of_first_n = noise_transitions * (noise_transitions + 1) / 2
    sum_of_first_n_squares = sum_of_first_n * (2 * noise_transitions + 1) / 3
    return _pack_and_reshape(
        num_steps * a + sum_of_first_n * (b + c) + sum_of_first_n_squares * d,
        num_steps * b + sum_of_first_n * d,
        num_steps * c + sum_of_first_n * d,
        num_steps * d)