from tensorflow.python.framework import tensor_shape


ops.RegisterShape("Softmax")(
    common_shapes.unchanged_shape_with_rank(2))


@ops.RegisterShape("InTopK")
def _InTopKShape(op):
  """Shape function for InTopK op."""
  predictions_shape = op.inputs[0].get_shape().with_rank(2)
  targets_shape = op.inputs[1].get_shape().with_rank(1)
  batch_size = predictions_shape[0].merge_with(targets_shape[0])
  return [tensor_shape.vector(batch_size.value)]


@ops.RegisterShape("TopK")
@ops.RegisterShape("TopKV2")
def _TopKShape(op):
  """Shape function for TopK and TopKV2 ops."""
  input_shape = op.inputs[0].get_shape().with_rank_at_least(1)
  if len(op.inputs) >= 2:
    k = tensor_util.constant_value(op.inputs[1])
  else:
