from tensorflow.python.framework import ops

  return [op.inputs[1].get_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())]
