from tensorflow.python.framework import ops

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


@ops.RegisterShape("ScatterAdd")
@ops.RegisterShape("ScatterSub")
@ops.RegisterShape("ScatterUpdate")
def _ScatterUpdateShape(op):
  """Shape function for the sparse update ops."""
  var_shape = op.inputs[0].get_shape()
  indices_shape = op.inputs[1].get_shape()
  unused_updates_shape = op.inputs[2].get_shape().merge_with(
      indices_shape.concatenate(var_shape[1:]))
  return [var_shape]
