def _get_sparse_tensors(self, inputs, weight_collections=None,
                         trainable=None):
   sparse_tensors = self.categorical_column._get_sparse_tensors(inputs)
   id_tensor = sparse_tensors.id_tensor
   weight_tensor = sparse_tensors.weight_tensor
   # Expands final dimension, so that embeddings are not combined during
   # embedding lookup.
   check_id_rank = check_ops.assert_equal(
       array_ops.rank(id_tensor), 2,
       data=[
           'Column {} expected ID tensor of rank 2. '.format(self.name),
           'id_tensor shape: ', array_ops.shape(id_tensor)])
   with ops.control_dependencies([check_id_rank]):
     id_tensor = sparse_ops.sparse_reshape(
         id_tensor,
         shape=array_ops.concat([id_tensor.dense_shape, [1]], axis=0))
   if weight_tensor is not None:
     check_weight_rank = check_ops.assert_equal(
         array_ops.rank(weight_tensor), 2,
         data=[
             'Column {} expected weight tensor of rank 2.'.format(self.name),
             'weight_tensor shape:', array_ops.shape(weight_tensor)])
     with ops.control_dependencies([check_weight_rank]):
       weight_tensor = sparse_ops.sparse_reshape(
           weight_tensor,
           shape=array_ops.concat([weight_tensor.dense_shape, [1]], axis=0))
   return fc._CategoricalColumn.IdWeightPair(id_tensor, weight_tensor)

 def testFeedMismatchedSizesWithInferredDim(self):
   with self.session(use_gpu=False) as sess:
     sp_input = self._SparseTensorPlaceholder()
     input_val = self._SparseTensorValue_5x6()
     sp_output = sparse_ops.sparse_reshape(sp_input, [4, -1])
     with self.assertRaisesOpError("requested shape requires a multiple"):
       sess.run(sp_output, {sp_input: input_val})

 def testFeedMultipleInferredDims(self):
   with self.session(use_gpu=False) as sess:
     sp_input = self._SparseTensorPlaceholder()
     input_val = self._SparseTensorValue_5x6()
     sp_output = sparse_ops.sparse_reshape(sp_input, [4, -1, -1])
     with self.assertRaisesOpError("only one output dimension may be -1"):
       sess.run(sp_output, {sp_input: input_val})

def dense_labels_to_sparse(dense, length):
  """Convert dense labels with sequence lengths to sparse tensor.

  Args:
    dense: tensor of shape [batch, max_length]
    length: int tensor of shape [batch]
      The length of each sequence in dense.

  Returns:
    tf.SparseTensor with values only for the valid elements of sequences.
  """

  flat_values = array_ops.reshape(dense, [-1])
  flat_indices = math_ops.range(
      array_ops.shape(flat_values, out_type=dtypes.int64)[0])
  mask = array_ops.sequence_mask(length, maxlen=array_ops.shape(dense)[1])
  flat_mask = array_ops.reshape(mask, [-1])
  indices = array_ops.expand_dims(
      array_ops.boolean_mask(flat_indices, flat_mask), 1)
  values = array_ops.boolean_mask(flat_values, flat_mask)
  sparse = sparse_tensor.SparseTensor(
      indices=indices, values=math_ops.cast(values, dtypes.int32),
      dense_shape=array_ops.shape(flat_values, out_type=dtypes.int64))
  reshaped = sparse_ops.sparse_reshape(sparse, array_ops.shape(dense))
  max_length = math_ops.reduce_max(length)
  return sparse_tensor.SparseTensor(
      indices=reshaped.indices,
      values=reshaped.values,
      dense_shape=[
          math_ops.cast(reshaped.dense_shape[0], dtypes.int64),
          math_ops.cast(max_length, dtypes.int64)])

  def testFeedDenseReshapeSemantics(self):
    with self.session(use_gpu=False) as sess:
      # Compute a random rank-5 initial shape and new shape, randomly sparsify
      # it, and check that the output of SparseReshape has the same semantics
      # as a dense reshape.
      factors = np.array([2] * 4 + [3] * 4 + [5] * 4)  # 810k total elements
      orig_rank = np.random.randint(2, 7)
      orig_map = np.random.randint(orig_rank, size=factors.shape)
      orig_shape = [np.prod(factors[orig_map == d]) for d in range(orig_rank)]
      new_rank = np.random.randint(2, 7)
      new_map = np.random.randint(new_rank, size=factors.shape)
      new_shape = [np.prod(factors[new_map == d]) for d in range(new_rank)]

      orig_dense = np.random.uniform(size=orig_shape)
      orig_indices = np.transpose(np.nonzero(orig_dense < 0.5))
      orig_values = orig_dense[orig_dense < 0.5]

      new_dense = np.reshape(orig_dense, new_shape)
      new_indices = np.transpose(np.nonzero(new_dense < 0.5))
      new_values = new_dense[new_dense < 0.5]

      sp_input = self._SparseTensorPlaceholder()
      input_val = sparse_tensor.SparseTensorValue(orig_indices, orig_values,
                                                  orig_shape)
      sp_output = sparse_ops.sparse_reshape(sp_input, new_shape)

      output_val = sess.run(sp_output, {sp_input: input_val})
      self.assertAllEqual(output_val.indices, new_indices)
      self.assertAllEqual(output_val.values, new_values)
      self.assertAllEqual(output_val.dense_shape, new_shape)

 def batch_reduce_fn(state, value):
   padded_value = sparse_tensor.SparseTensor(
       indices=value.indices, values=value.values, dense_shape=padded_shape)
   reshaped_value = sparse_ops.sparse_reshape(
       padded_value,
       array_ops.concat(
           [np.array([1], dtype=np.int64), padded_value.dense_shape], 0))
   return sparse_ops.sparse_concat(0, [state, reshaped_value])

 def testFeedMismatchedSizes(self):
   with self.session(use_gpu=False) as sess:
     sp_input = self._SparseTensorPlaceholder()
     input_val = self._SparseTensorValue_5x6()
     sp_output = sparse_ops.sparse_reshape(sp_input, [4, 7])
     with self.assertRaisesOpError(
         "Input to reshape is a tensor with 30 dense values"):
       sess.run(sp_output, {sp_input: input_val})

  def testSameShape(self):
    with self.session(use_gpu=False) as sess:
      input_val = self._SparseTensorValue_5x6()
      sp_output = sparse_ops.sparse_reshape(input_val, [5, 6])

      output_val = sess.run(sp_output)
      self.assertAllEqual(output_val.indices, input_val.indices)
      self.assertAllEqual(output_val.values, input_val.values)
      self.assertAllEqual(output_val.dense_shape, input_val.dense_shape)

  def testFeedSameShape(self):
    with self.session(use_gpu=False) as sess:
      sp_input = self._SparseTensorPlaceholder()
      input_val = self._SparseTensorValue_5x6()
      sp_output = sparse_ops.sparse_reshape(sp_input, [5, 6])

      output_val = sess.run(sp_output, {sp_input: input_val})
      self.assertAllEqual(output_val.indices, input_val.indices)
      self.assertAllEqual(output_val.values, input_val.values)
      self.assertAllEqual(output_val.dense_shape, input_val.dense_shape)

  def testWorksWellWithTfShape(self):
    with self.session(use_gpu=False) as sess:
      sp_input = self._SparseTensorPlaceholder()
      input_val = self._SparseTensorValue_5x6()
      shape = array_ops.shape(sp_input)  # tf.shape generates int32 output
      sp_output = sparse_ops.sparse_reshape(sp_input, shape)

      output_val = sess.run(sp_output, {sp_input: input_val})
      self.assertAllEqual(output_val.indices, input_val.indices)
      self.assertAllEqual(output_val.values, input_val.values)
      self.assertAllEqual(output_val.dense_shape, input_val.dense_shape)

  def testUpRank(self):
    with self.session(use_gpu=False) as sess:
      input_val = self._SparseTensorValue_5x6()
      sp_output = sparse_ops.sparse_reshape(input_val, [2, 3, 5])

      output_val = sess.run(sp_output)
      self.assertAllEqual(output_val.indices,
                          np.array([[0, 0, 0], [0, 1, 1], [0, 1, 4], [0, 2, 0],
                                    [1, 1, 0], [1, 1, 1]]))
      self.assertAllEqual(output_val.values, input_val.values)
      self.assertAllEqual(output_val.dense_shape, [2, 3, 5])

  def testFeedUpRankWithInferredDim(self):
    with self.session(use_gpu=False) as sess:
      sp_input = self._SparseTensorPlaceholder()
      input_val = self._SparseTensorValue_5x6()
      sp_output = sparse_ops.sparse_reshape(sp_input, [2, -1, 5])

      output_val = sess.run(sp_output, {sp_input: input_val})
      self.assertAllEqual(output_val.indices,
                          np.array([[0, 0, 0], [0, 1, 1], [0, 1, 4], [0, 2, 0],
                                    [1, 1, 0], [1, 1, 1]]))
      self.assertAllEqual(output_val.values, input_val.values)
      self.assertAllEqual(output_val.dense_shape, [2, 3, 5])

  def testFeedNewShapeSameRank(self):
    with self.session(use_gpu=False) as sess:
      sp_input = self._SparseTensorPlaceholder()
      input_val = self._SparseTensorValue_5x6()
      sp_output = sparse_ops.sparse_reshape(sp_input, [3, 10])

      output_val = sess.run(sp_output, {sp_input: input_val})
      self.assertAllEqual(output_val.indices,
                          np.array([[0, 0], [0, 6], [0, 9], [1, 0], [2, 0],
                                    [2, 1]]))
      self.assertAllEqual(output_val.values, input_val.values)
      self.assertAllEqual(output_val.dense_shape, [3, 10])

  def testFeedDownRankWithInferredDim(self):
    with self.session(use_gpu=False) as sess:
      sp_input = self._SparseTensorPlaceholder()
      input_val = self._SparseTensorValue_2x3x4()
      sp_output = sparse_ops.sparse_reshape(sp_input, [6, -1])

      output_val = sess.run(sp_output, {sp_input: input_val})
      self.assertAllEqual(output_val.indices,
                          np.array([[0, 1], [1, 0], [1, 2], [3, 3], [4, 1],
                                    [4, 3], [5, 2]]))
      self.assertAllEqual(output_val.values, input_val.values)
      self.assertAllEqual(output_val.dense_shape, [6, 4])

def _maybe_reshape_input_tensor(tensor, column_name, output_rank):
  """Reshape the input tensor by the following rule.

  1. If `output_rank > input_rank + 1`, raise a `ValueError`.
  2. If `output_rank == input_rank + 1`, expand the tensor by one dimension.
  3. If `output_rank == input_rank`, do nothing.
  4. If `output_rank < input_rank`, flatten the inner dimensions of the tensor.

  Args:
    tensor: A Tensor or SparseTensor to be reshaped.
    column_name: A string name of the feature column for the tensor.
    output_rank: the desired rank of the tensor.
  Returns:
    A reshaped Tensor or SparseTensor.
  Raises:
    ValueError: if `output_rank > input_rank + 1` for the input tensor.
  """
  input_rank = tensor.get_shape().ndims

  if input_rank is None and isinstance(tensor, sparse_tensor_py.SparseTensor):
    # Try to get the rank of a sparse tensor by its dense_shape's shape.
    input_rank = tensor.dense_shape.get_shape().as_list()[0]

  if input_rank is None:
    raise ValueError('Error while processing column {}. Rank of input Tensor '
                     'can not be None.'.format(column_name))

  if output_rank > input_rank + 1:
    raise ValueError('Error while processing column {}. Rank of input Tensor '
                     '({}) should be the same as output_rank ({}). For '
                     'example, sequence data should typically be 3 '
                     'dimensional (rank 3) while non-sequence data is '
                     'typically 2 dimensional (rank 2).'.format(
                         column_name, input_rank, output_rank))
  elif output_rank == input_rank + 1:
    # Expand the tensor's shape by 1 dimension.
    if isinstance(tensor, sparse_tensor_py.SparseTensor):
      output_shape = array_ops.concat([tensor.dense_shape, [1]], 0)
      return sparse_ops.sparse_reshape(tensor, output_shape)
    else:
      reshaped = array_ops.expand_dims(tensor, -1)
      # Try to calculate the new shape.
      static_shape = tensor.get_shape()
      if static_shape is not None and static_shape.dims is not None:
        reshaped.set_shape(static_shape.as_list() + [1])
      return reshaped
  elif output_rank < input_rank:
    return layers._inner_flatten(tensor, output_rank)  # pylint: disable=protected-access
  else:
    return tensor

  def testFeedPartialShapes(self):
    with self.session(use_gpu=False):
      # Incorporate new rank into shape information if known
      sp_input = self._SparseTensorPlaceholder()
      sp_output = sparse_ops.sparse_reshape(sp_input, [2, 3, 5])
      self.assertListEqual(sp_output.indices.get_shape().as_list(), [None, 3])
      self.assertListEqual(sp_output.dense_shape.get_shape().as_list(), [3])

      # Incorporate known shape information about input indices in output
      # indices
      sp_input = self._SparseTensorPlaceholder()
      sp_input.indices.set_shape([5, None])
      sp_output = sparse_ops.sparse_reshape(sp_input, [2, 3, 5])
      self.assertListEqual(sp_output.indices.get_shape().as_list(), [5, 3])
      self.assertListEqual(sp_output.dense_shape.get_shape().as_list(), [3])

      # Even if new_shape has no shape information, we know the ranks of
      # output indices and shape
      sp_input = self._SparseTensorPlaceholder()
      sp_input.indices.set_shape([5, None])
      new_shape = array_ops.placeholder(dtypes.int64)
      sp_output = sparse_ops.sparse_reshape(sp_input, new_shape)
      self.assertListEqual(sp_output.indices.get_shape().as_list(), [5, None])
      self.assertListEqual(sp_output.dense_shape.get_shape().as_list(), [None])

 def tensors_to_item(self, keys_to_tensors):
     tensor = keys_to_tensors[self._tensor_key]
     shape = self._shape
     if self._shape_key:
         shape = keys_to_tensors[self._shape_key]
         if isinstance(shape, ops.SparseTensor):
             shape = sparse_ops.sparse_tensor_to_dense(shape)
     if isinstance(tensor, ops.SparseTensor):
         if shape is not None:
             tensor = sparse_ops.sparse_reshape(tensor, shape)
         tensor = sparse_ops.sparse_tensor_to_dense(tensor, self._default_value)
     else:
         if shape is not None:
             tensor = array_ops.reshape(tensor, shape)
     return tensor

 def tensors_to_item(self, keys_to_tensors):
   tensor = keys_to_tensors[self._tensor_key]
   shape = self._shape
   if self._shape_keys:
     shape_dims = []
     for k in self._shape_keys:
       shape_dim = keys_to_tensors[k]
       if isinstance(shape_dim, ops.SparseTensor):
         shape_dim = sparse_ops.sparse_tensor_to_dense(shape_dim)
       shape_dims.append(shape_dim)
     shape = array_ops.squeeze(array_ops.pack(shape_dims))
   if isinstance(tensor, ops.SparseTensor):
     if shape is not None:
       tensor = sparse_ops.sparse_reshape(tensor, shape)
     tensor = sparse_ops.sparse_tensor_to_dense(tensor, self._default_value)
   else:
     if shape is not None:
       tensor = array_ops.reshape(tensor, shape)
   return tensor

 def _structuredRaggedSparseElement(self, structure, shapes, dtype,
                                    padded_shape):
   if structure is None:
     dense_shape = np.maximum(np.amax(shapes, axis=0), padded_shape)
     values = []
     for shape in shapes:
       dense_to_sparse = self._make_dense_to_sparse_fn(len(shape) == 0)  # pylint: disable=g-explicit-length-test
       sparse = dense_to_sparse(array_ops.zeros(shape, dtype=dtype))
       padded_sparse = sparse_tensor.SparseTensor(sparse.indices,
                                                  sparse.values, dense_shape)
       reshaped_sparse = sparse_ops.sparse_reshape(
           padded_sparse,
           array_ops.concat([np.array([1], dtype=np.int64), dense_shape], 0))
       values.append(reshaped_sparse)
     return sparse_ops.sparse_concat(0, values)
   else:
     return tuple([
         self._structuredRaggedSparseElement(substructure, shapes, dtype,
                                             padded_shape)
         for substructure in structure
     ])

 def testRaisesIfMoreThanOneInferredDim(self):
   sp_input = sparse_tensor.SparseTensor.from_value(
       self._SparseTensorValue_2x3x4())
   with self.assertRaisesRegexp(ValueError, "At most one dimension can"):
     sparse_ops.sparse_reshape(sp_input, shape=(-1, 2, -1))