Register operators with different tensor and scalar versions. If `clazz_object` is `SparseTensor`, assumes `func` takes `(sp_indices, sp_values, sp_shape, dense)` and outputs `(new_sp_values)`. Args: func: the operator op_name: name of the operator being overridden clazz_object:
(func, op_name, clazz_object=ops.Tensor)
| 882 | |
| 883 | |
| 884 | def _OverrideBinaryOperatorHelper(func, op_name, clazz_object=ops.Tensor): |
| 885 | """Register operators with different tensor and scalar versions. |
| 886 | |
| 887 | If `clazz_object` is `SparseTensor`, assumes `func` takes `(sp_indices, |
| 888 | sp_values, sp_shape, dense)` and outputs `(new_sp_values)`. |
| 889 | |
| 890 | Args: |
| 891 | func: the operator |
| 892 | op_name: name of the operator being overridden |
| 893 | clazz_object: class to override for. Either `Tensor` or `SparseTensor`. |
| 894 | """ |
| 895 | |
| 896 | def binary_op_wrapper(x, y): |
| 897 | with ops.name_scope(None, op_name, [x, y]) as name: |
| 898 | if isinstance(x, ops.Tensor) and isinstance(y, ops.Tensor): |
| 899 | return func(x, y, name=name) |
| 900 | elif not isinstance(y, sparse_tensor.SparseTensor): |
| 901 | try: |
| 902 | y = ops.convert_to_tensor_v2( |
| 903 | y, dtype_hint=x.dtype.base_dtype, name="y") |
| 904 | except TypeError: |
| 905 | # If the RHS is not a tensor, it might be a tensor aware object |
| 906 | # that can implement the operator with knowledge of itself |
| 907 | # and the tensor. |
| 908 | if hasattr(type(y), "__r%s__" % op_name): |
| 909 | return NotImplemented |
| 910 | else: |
| 911 | raise |
| 912 | return func(x, y, name=name) |
| 913 | |
| 914 | def binary_op_wrapper_sparse(sp_x, y): |
| 915 | with ops.name_scope(None, op_name, [sp_x, y]) as name: |
| 916 | y = ops.convert_to_tensor(y, dtype=sp_x.dtype.base_dtype, name="y") |
| 917 | return sparse_tensor.SparseTensor( |
| 918 | sp_x.indices, |
| 919 | func(sp_x.indices, sp_x.values, sp_x.dense_shape, y, name=name), |
| 920 | sp_x.dense_shape) |
| 921 | |
| 922 | def r_binary_op_wrapper(y, x): |
| 923 | with ops.name_scope(None, op_name, [x, y]) as name: |
| 924 | x = ops.convert_to_tensor(x, dtype=y.dtype.base_dtype, name="x") |
| 925 | return func(x, y, name=name) |
| 926 | |
| 927 | # Propagate func.__doc__ to the wrappers |
| 928 | try: |
| 929 | doc = func.__doc__ |
| 930 | except AttributeError: |
| 931 | doc = None |
| 932 | binary_op_wrapper.__doc__ = doc |
| 933 | r_binary_op_wrapper.__doc__ = doc |
| 934 | binary_op_wrapper_sparse.__doc__ = doc |
| 935 | |
| 936 | if clazz_object is ops.Tensor: |
| 937 | clazz_object._override_operator("__%s__" % op_name, binary_op_wrapper) |
| 938 | del binary_op_wrapper |
| 939 | clazz_object._override_operator("__r%s__" % op_name, r_binary_op_wrapper) |
| 940 | del r_binary_op_wrapper |
| 941 | else: |
no test coverage detected