from tensorflow.python.framework import tensor_util

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:
    k = op.get_attr("k")
  last = input_shape[-1].value
