from tensorflow.python.framework import tensor_shape

@ops.RegisterShape("AssignAdd")
@ops.RegisterShape("AssignSub")
def _AssignUpdateShape(op):
  """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())]


@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()
