from tensorflow.python.framework import ops


def _OverrideBinaryOperatorHelper(func, op_name):
  """Register operators with different tensor and scalar versions.

  Args:
    func: the operator
    op_name: name of the operator being overridden
  """

  def binary_op_wrapper(x, y):
    with ops.op_scope([x, y], None, op_name) as name:
      assert isinstance(x, ops.Tensor)
      y = ops.convert_to_tensor(y, dtype=x.dtype.base_dtype, name="y")
      return func(x, y, name=name)

  ops.Tensor._override_operator("__%s__" % op_name, binary_op_wrapper)
  del binary_op_wrapper

  def r_binary_op_wrapper(y, x):
    with ops.op_scope([x, y], None, op_name) as name:
