from tensorflow.python.framework import ops

  """Shape function for the AssignAdd and AssignSub dense update ops."""
  return [op.inputs[0].get_shape().merge_with(op.inputs[1].get_shape())]


@ops.RegisterShape("CountUpTo")
def _CountUpToShape(op):
  """Shape function for the CountUpTo op."""
  return [op.inputs[0].get_shape().merge_with(tensor_shape.scalar())]
