from tensorflow.python.framework import ops

@ops.RegisterShape("Complex")
@ops.RegisterShape("Div")
@ops.RegisterShape("Equal")
@ops.RegisterShape("Greater")
@ops.RegisterShape("GreaterEqual")
@ops.RegisterShape("Less")
@ops.RegisterShape("LessEqual")
@ops.RegisterShape("LogicalAnd")
@ops.RegisterShape("LogicalOr")
@ops.RegisterShape("Maximum")
@ops.RegisterShape("Minimum")
@ops.RegisterShape("Mod")
@ops.RegisterShape("Mul")
@ops.RegisterShape("NotEqual")
@ops.RegisterShape("Pow")
@ops.RegisterShape("Sub")
def _BroadcastShape(op):
  """Common shape function for binary operators that broadcast their inputs."""
  shape_x = op.inputs[0].get_shape()
  shape_y = op.inputs[1].get_shape()
  if shape_x.ndims is None or shape_y.ndims is None:
    return [tensor_shape.unknown_shape()]

  # To compute the broadcasted dimensions, we zip together shape_x and shape_y,
