subtraction for tf sparsetensors either x1 or x2 or both can be tf sparsetensors args x1 fist tensor to add x2 second tensor to add returns the sum of x1 and x2 which is a tf sparsetensor if and only if both x1 or x2 are tf sparsetensors compute the indices for the union of the indices of the provided tf sparsetensors and another set of indices and return the modified values for these indices args x a tf sparsetensor indices another set of indices in the tf sparsetensor format returns a tuple containing the indices for the union x1 values for the union indices some zeros were added x2 values for the union indices some zeros were added or none if x2values was none add zeros at the x2 indices to x1 to create the union add zeros at the x1 indices to x2 to create the union compute the indices for the union of two tf indexedslices and modify the values for these indices args x1 the first tf indexedslices x2indices the indices for the second tf indexedslices x2value optional the values for the second tf indexedslices returns a tuple containing the indices for the union x1 values for the union indices some zeros were added x2 values for the union indices some zeros were added or none if x2values was none compute the union of the indices by doing a logical or between the onehot encoded indices for x1 and x2 regather the values with extra zeros added at indices that are part of the union but were not in x1 or x2 only recompute values if some indices were added compute the indices for the intersection of two tf sparsetensors and modify the values for these indices args x1 the first tf sparsetensor x2 the second tf sparsetensor returns a tuple containing the indices for the intersection x1 values for the intersection indices some values were removed x2 values for the intersection indices some values were removed compute the intersection of indices in the form of a sparse tensor containing ones as values tf sets intersection ignores the last dimension when so we need to add a dummy extra dimension and then remove it compute the masks to remove indices in x1 and x2 that are not in the intersection then trim x1 and x2 compute the indices for the intersection of two tf indexedslices and modify the values for these indices args x1 the first tf indexedslices x2 the second tf indexedslices returns a tuple containing the indices for the intersection x1 values for the intersection indices some values were removed x2 values for the intersection indices some values were removed compute the intersection of the indices by doing a logical and between the one hot encoded indices for x1 and x2 regather sub parts of the values that are part of the intersection only recompute values if some indices were removed decorator to add support for tf sparsetensor and tf indexedslices to a nonzeropreserving elementwise unary operator there are requirements on the operator for this decorator to work correctly the operator must be elementwise the operator must be unary one input tensor and one output tensor the operator must return a tensor of the same shape additional arguments to the function besides the input tensor are supported the returned result is a dense tensor and contains defaultvalue outside of the indices of the input tensor args defaultvalue the value to use outside of indices it must be the value that the operator returns for zero values returns wrapped function that supports tf sparsetensor and tf indexedslices decorator to add support for tf sparsetensor and tf indexedslices to a zeropreserving elementwise unary operator there are requirements on the operator for this decorator to work correctly the operator must be elementwise the operator must be unary one input tensor and one output tensor the operator must return a tensor of the same shape and if it is a tf sparsetensor or tf indexedslices the indices of the result must be the same therefore reduction operations are not supported e g mean operation for which the result may be dense e g reciprocal or the sparse indices depend on the inputs are not supported e g clip this implies that func0 must be 0 additional arguments to the function besides the input tensor are supported as long as they cannot change the indices of the result for instance round is supported but clip is not supported as clipx 1 0 2 0 would always return a dense tensor note that if an input sparse tensor contains zero values the indices and the zero values are preserved args func the function to wrap returns wrapped function that supports tf sparsetensor and tf indexedslices decorator to add support for tf sparsetensor and tf indexedslices to an elementwise binary operator such that the indices present in the result are the union of the indices in the two operand the primary use case for this is the add and subtract operators there are requirements on the operator for this decorator to work correctly the operator must be elementwise the operator must be binary two input tensors and one output tensor both inputs must be of the same shape or one input must be a scalar the output must be of the same shape as the non scalar inputs the indices of the output must be the union of the indices of the inputs this implies that func0 0 must be 0 as a result if one operand is dense or a scalar then the result will be dense additional arguments to the function besides the input tensors are not supported note that if the result of the operation is zero at some indices including because the operands were zero at these indices the zeros and indices are preserved args sparseop implementation of the operation for tf sparsetensor must work if both of the operands are tf sparsetensors and can optionally work if one of the operand is a tf sparsetensor and the other one is dense tensor see densifymixed densifymixed if true sparseop does not support a mix of tf sparsetensor and dense tensor or dense tensor with tf sparsetensor and the tf sparsetensor tensor is densified returns wrapped function that supports tf sparsetensor and tf indexedslices x1 is a sparsetensor and x2 is a sparsetensor x1 is a sparsetensor x2 is a scalar broadcast x2 is a sparsetensor x1 is a scalar broadcast x1 is an indexedslices and x2 is an indexedslices compute the union of indices now it is an elementwise operation on the union x1 is an indexedslices densify x2 is an indexedslices densify decorator to add support for tf sparsetensor and tf indexedslices to an elementwise binary operator such that the indices present in the result are the intersection of the indices in the two operand the primary use case for this is the multiply operator there are requirements on the operator for this decorator to work correctly the operator must be elementwise the operator must be binary two input tensors and one output tensor both inputs must be of the same shape or one input must be a scalar the output must be of the same shape as the non scalar inputs the indices of the output must be the intersection of the indices of the inputs this implies that func0 x and funcx 0 must be 0 for any x as a result if one operand is dense or a scalar then the indices are the ones from the other operand additional arguments to the function besides the input tensors are not supported note that if the operands contains zero values at some common indices the indices and the zero values are preserved args func the function to wrap returns wrapped function that supports tf sparsetensor and tf indexedslices x1 is a sparsetensor and x2 is a sparsetensor compute the intersection of indices now it is an elementwise operation on the intersection x1 is a sparsetensor x2 is a scalar apply func elementwise x2 is dense gather values from x1 indices x2 is a sparsetensor x1 is a scalar apply func elementwise x1 is dense gather values from x2 indices x1 is an indexedslices and x2 is an indexedslices compute the intersection of indices now it is an elementwise operation on the intersection x1 is an indexedslices x2 is a scalar apply func elementwise x2 is dense gather values from x1 indices x2 is an indexedslices x1 is a scalar apply func elementwise x1 is dense gather values from x2 indices default case no sparsetensor and no indexedslices decorator to add support for tf sparsetensor and tf indexedslices to elementwise binary division and related operators this decorator is designed for operations related to the division of the two operands e g divide mod it accepts tf sparsetensor and tf indexedslices for both the dividend and the divisor but handles them differently based on whether they are the dividend or the divisor if the divisor is a tf sparsetensor or tf indexedslices it is densified and the result is dense because the result contains inf or nan outside of the indices of the dividend if the dividend is a tf sparsetensor or tf indexedslices and the divisor is dense it finds occurrences of zeros and nans in the divisor the result may therefore have more indices than there were in the dividend to return correct values where the divisor was zero or nan if the dividend is a tf sparsetensor or tf indexedslices and the divisor is a scalar it does the division elementwise note that the result is incorrectly sparse if the scalar divisor is zero args func the function to wrap returns wrapped function that supports tf sparsetensor and tf indexedslices x1 is a sparsetensor and x2 is a sparsetensor divisor is sparse meaning we re doing divisions by zero outside of x2 indices so the result is dense densify both x1 is a sparsetensor x2 is a scalar apply func elementwise x2 is dense gather values from x1 indices compute the union of indices to keep zeros and nans x2 is a sparsetensor divisor is sparse densify to do the divisions by zero correctly x1 is an indexedslices and x2 is an indexedslices divisor is slices meaning we re doing divisions by zero outside of x2 indices so the result is dense densify both x1 is a indexedslices x2 is a scalar apply func elementwise x2 is dense gather values from x1 indices compute the union of indices to keep zeros and nans x2 is a indexedslices divisor is slices densify to do the divisions by zero correctly default case no sparsetensor and no indexedslices subtraction for tf sparsetensor s either x1 or x2 or both can be tf sparsetensor s args x1 fist tensor to add x2 second tensor to add returns the sum of x1 and x2 which is a tf sparsetensor if and only if both x1 or x2 are tf sparsetensor s compute the indices for the union of the indices of the provided tf sparsetensor s and another set of indices and return the modified values for these indices args x a tf sparsetensor indices another set of indices in the tf sparsetensor format returns a tuple containing the indices for the union x1 values for the union indices some zeros were added x2 values for the union indices some zeros were added or none if x2_values was none add zeros at the x2 indices to x1 to create the union add zeros at the x1 indices to x2 to create the union compute the indices for the union of two tf indexedslices and modify the values for these indices args x1 the first tf indexedslices x2_indices the indices for the second tf indexedslices x2_value optional the values for the second tf indexedslices returns a tuple containing the indices for the union x1 values for the union indices some zeros were added x2 values for the union indices some zeros were added or none if x2_values was none compute the union of the indices by doing a logical or between the one hot encoded indices for x1 and x2 re gather the values with extra zeros added at indices that are part of the union but were not in x1 or x2 only recompute values if some indices were added compute the indices for the intersection of two tf sparsetensor s and modify the values for these indices args x1 the first tf sparsetensor x2 the second tf sparsetensor returns a tuple containing the indices for the intersection x1 values for the intersection indices some values were removed x2 values for the intersection indices some values were removed compute the intersection of indices in the form of a sparse tensor containing ones as values tf sets intersection ignores the last dimension when so we need to add a dummy extra dimension and then remove it compute the masks to remove indices in x1 and x2 that are not in the intersection then trim x1 and x2 compute the indices for the intersection of two tf indexedslices and modify the values for these indices args x1 the first tf indexedslices x2 the second tf indexedslices returns a tuple containing the indices for the intersection x1 values for the intersection indices some values were removed x2 values for the intersection indices some values were removed compute the intersection of the indices by doing a logical and between the one hot encoded indices for x1 and x2 re gather sub parts of the values that are part of the intersection only recompute values if some indices were removed decorator to add support for tf sparsetensor and tf indexedslices to a non zero preserving element wise unary operator there are requirements on the operator for this decorator to work correctly the operator must be element wise the operator must be unary one input tensor and one output tensor the operator must return a tensor of the same shape additional arguments to the function besides the input tensor are supported the returned result is a dense tensor and contains default_value outside of the indices of the input tensor args default_value the value to use outside of indices it must be the value that the operator returns for zero values returns wrapped function that supports tf sparsetensor and tf indexedslices decorator to add support for tf sparsetensor and tf indexedslices to a zero preserving element wise unary operator there are requirements on the operator for this decorator to work correctly the operator must be element wise the operator must be unary one input tensor and one output tensor the operator must return a tensor of the same shape and if it is a tf sparsetensor or tf indexedslices the indices of the result must be the same therefore reduction operations are not supported e g mean operation for which the result may be dense e g reciprocal or the sparse indices depend on the inputs are not supported e g clip this implies that func 0 must be 0 additional arguments to the function besides the input tensor are supported as long as they cannot change the indices of the result for instance round is supported but clip is not supported as clip x 1 0 2 0 would always return a dense tensor note that if an input sparse tensor contains zero values the indices and the zero values are preserved args func the function to wrap returns wrapped function that supports tf sparsetensor and tf indexedslices decorator to add support for tf sparsetensor and tf indexedslices to an element wise binary operator such that the indices present in the result are the union of the indices in the two operand the primary use case for this is the add and subtract operators there are requirements on the operator for this decorator to work correctly the operator must be element wise the operator must be binary two input tensors and one output tensor both inputs must be of the same shape or one input must be a scalar the output must be of the same shape as the non scalar inputs the indices of the output must be the union of the indices of the inputs this implies that func 0 0 must be 0 as a result if one operand is dense or a scalar then the result will be dense additional arguments to the function besides the input tensors are not supported note that if the result of the operation is zero at some indices including because the operands were zero at these indices the zeros and indices are preserved args sparse_op implementation of the operation for tf sparsetensor must work if both of the operands are tf sparsetensor s and can optionally work if one of the operand is a tf sparsetensor and the other one is dense tensor see densify_mixed densify_mixed if true sparse_op does not support a mix of tf sparsetensor and dense tensor or dense tensor with tf sparsetensor and the tf sparsetensor tensor is densified returns wrapped function that supports tf sparsetensor and tf indexedslices x1 is a sparsetensor and x2 is a sparsetensor x1 is a sparsetensor x2 is a scalar broadcast x2 is a sparsetensor x1 is a scalar broadcast x1 is an indexedslices and x2 is an indexedslices compute the union of indices now it is an element wise operation on the union x1 is an indexedslices densify x2 is an indexedslices densify decorator to add support for tf sparsetensor and tf indexedslices to an element wise binary operator such that the indices present in the result are the intersection of the indices in the two operand the primary use case for this is the multiply operator there are requirements on the operator for this decorator to work correctly the operator must be element wise the operator must be binary two input tensors and one output tensor both inputs must be of the same shape or one input must be a scalar the output must be of the same shape as the non scalar inputs the indices of the output must be the intersection of the indices of the inputs this implies that func 0 x and func x 0 must be 0 for any x as a result if one operand is dense or a scalar then the indices are the ones from the other operand additional arguments to the function besides the input tensors are not supported note that if the operands contains zero values at some common indices the indices and the zero values are preserved args func the function to wrap returns wrapped function that supports tf sparsetensor and tf indexedslices x1 is a sparsetensor and x2 is a sparsetensor compute the intersection of indices now it is an element wise operation on the intersection x1 is a sparsetensor x2 is a scalar apply func element wise x2 is dense gather values from x1 indices x2 is a sparsetensor x1 is a scalar apply func element wise x1 is dense gather values from x2 indices x1 is an indexedslices and x2 is an indexedslices compute the intersection of indices now it is an element wise operation on the intersection x1 is an indexedslices x2 is a scalar apply func element wise x2 is dense gather values from x1 indices x2 is an indexedslices x1 is a scalar apply func element wise x1 is dense gather values from x2 indices default case no sparsetensor and no indexedslices decorator to add support for tf sparsetensor and tf indexedslices to element wise binary division and related operators this decorator is designed for operations related to the division of the two operands e g divide mod it accepts tf sparsetensor and tf indexedslices for both the dividend and the divisor but handles them differently based on whether they are the dividend or the divisor if the divisor is a tf sparsetensor or tf indexedslices it is densified and the result is dense because the result contains inf or nan outside of the indices of the dividend if the dividend is a tf sparsetensor or tf indexedslices and the divisor is dense it finds occurrences of zeros and nans in the divisor the result may therefore have more indices than there were in the dividend to return correct values where the divisor was zero or nan if the dividend is a tf sparsetensor or tf indexedslices and the divisor is a scalar it does the division element wise note that the result is incorrectly sparse if the scalar divisor is zero args func the function to wrap returns wrapped function that supports tf sparsetensor and tf indexedslices x1 is a sparsetensor and x2 is a sparsetensor divisor is sparse meaning we re doing divisions by zero outside of x2 indices so the result is dense densify both x1 is a sparsetensor x2 is a scalar apply func element wise x2 is dense gather values from x1 indices compute the union of indices to keep zeros and nans x2 is a sparsetensor divisor is sparse densify to do the divisions by zero correctly x1 is an indexedslices and x2 is an indexedslices divisor is slices meaning we re doing divisions by zero outside of x2 indices so the result is dense densify both x1 is a indexedslices x2 is a scalar apply func element wise x2 is dense gather values from x1 indices compute the union of indices to keep zeros and nans x2 is a indexedslices divisor is slices densify to do the divisions by zero correctly default case no sparsetensor and no indexedslices
import functools

import tensorflow as tf

ones_bool = functools.partial(tf.ones, dtype=tf.bool)
ones_int8 = functools.partial(tf.ones, dtype=tf.int8)
zeros_int8 = functools.partial(tf.zeros, dtype=tf.int8)
ones_like_int8 = functools.partial(tf.ones_like, dtype=tf.int8)
zeros_like_int8 = functools.partial(tf.zeros_like, dtype=tf.int8)


def empty_tensor(shape, dtype):
    return tf.reshape(tf.convert_to_tensor((), dtype=dtype), shape=shape)


def sparse_to_dense(x, default_value=None):
    x_shape = x.shape
    x = tf.sparse.to_dense(x, default_value=default_value)
    x.set_shape(x_shape)
    return x


def sparse_with_values(x, values):
    x_shape = x.shape
    x = tf.SparseTensor(x.indices, values, x.dense_shape)
    x.set_shape(x_shape)
    return x


def broadcast_scalar_to_sparse_shape(scalar, sparse):
    output = tf.broadcast_to(scalar, sparse.dense_shape)
    output.set_shape(sparse.shape)
    return output


def sparse_subtract(x1, x2):
    
    if isinstance(x2, tf.SparseTensor):
        return tf.sparse.add(x1, tf.sparse.map_values(tf.negative, x2))
    else:
        return tf.sparse.add(x1, tf.negative(x2))


def sparse_union_indices_and_values(x1, x2_indices, x2_values=None):
    
    
    zeros2 = tf.SparseTensor(
        x2_indices,
        tf.zeros((tf.shape(x2_indices)[0],), x1.values.dtype),
        x1.dense_shape,
    )
    x1_for_union = tf.sparse.add(x1, zeros2)
    if x2_values is not None:
        
        x2 = tf.SparseTensor(x2_indices, x2_values, x1.dense_shape)
        zeros1 = tf.sparse.map_values(tf.zeros_like, x1)
        x2_for_union = tf.sparse.add(x2, zeros1)
        return x1_for_union.indices, x1_for_union.values, x2_for_union.values
    else:
        return x1_for_union.indices, x1_for_union.values, None


def indexed_slices_union_indices_and_values(x1, x2_indices, x2_values=None):
    
    
    
    dim_0 = x1.dense_shape[0]
    x1_indices_expanded = tf.expand_dims(x1.indices, axis=1)
    x2_indices_expanded = tf.expand_dims(x2_indices, axis=1)
    x1_indices_count = tf.shape(x1_indices_expanded)[0]
    x2_indices_count = tf.shape(x2_indices_expanded)[0]
    x1_indices_one_hot = tf.scatter_nd(
        x1_indices_expanded,
        ones_bool((x1_indices_count,)),
        (dim_0,),
    )
    x2_indices_one_hot = tf.scatter_nd(
        x2_indices_expanded,
        ones_bool((x2_indices_count,)),
        (dim_0,),
    )
    union_indices = tf.squeeze(
        tf.where(tf.math.logical_or(x1_indices_one_hot, x2_indices_one_hot)),
        axis=-1,
    )
    union_indices_count = tf.shape(union_indices)[0]

    
    
    def values_for_union(indices_expanded, indices_count, values):
        indices_indices = tf.scatter_nd(
            indices_expanded,
            tf.range(1, indices_count + 1),
            (dim_0,),
        )
        to_union_indices = tf.gather(indices_indices, union_indices)
        values_with_leading_zeros = tf.concat(
            [tf.zeros((1,) + values.shape[1:], values.dtype), values], axis=0
        )
        return tf.gather(values_with_leading_zeros, to_union_indices)

    
    x1_values_for_union_indices = tf.cond(
        tf.equal(x1_indices_count, union_indices_count),
        lambda: x1.values,
        lambda: values_for_union(
            x1_indices_expanded, x1_indices_count, x1.values
        ),
    )
    if x2_values is not None:
        x2_values_for_union_indices = tf.cond(
            tf.equal(x2_indices_count, union_indices_count),
            lambda: x2_values,
            lambda: values_for_union(
                x2_indices_expanded, x2_indices_count, x2_values
            ),
        )
    else:
        x2_values_for_union_indices = None

    return (
        union_indices,
        x1_values_for_union_indices,
        x2_values_for_union_indices,
    )


def sparse_intersection_indices_and_values(x1, x2):
    
    
    
    ones1 = tf.sparse.map_values(ones_like_int8, x1)
    ones2 = tf.sparse.map_values(ones_like_int8, x2)
    
    
    intersection_extra_dim = tf.sets.intersection(
        tf.sparse.expand_dims(ones1, axis=-1),
        tf.sparse.expand_dims(ones2, axis=-1),
    )

    def empty_intersection():
        return (
            empty_tensor((0, x1.shape.rank), dtype=tf.int64),
            empty_tensor((0,), dtype=x1.values.dtype),
            empty_tensor((0,), dtype=x2.values.dtype),
        )

    def non_empty_intersection():
        intersection = tf.sparse.reshape(intersection_extra_dim, x1.dense_shape)

        
        
        zeros1 = tf.sparse.map_values(zeros_like_int8, x1)
        zeros2 = tf.sparse.map_values(zeros_like_int8, x2)
        mask1 = tf.sparse.add(zeros1, intersection)
        mask2 = tf.sparse.add(zeros2, intersection)
        return (
            intersection.indices,
            tf.sparse.retain(x1, tf.cast(mask1.values, tf.bool)).values,
            tf.sparse.retain(x2, tf.cast(mask2.values, tf.bool)).values,
        )

    return tf.cond(
        tf.equal(tf.size(intersection_extra_dim), 0),
        empty_intersection,
        non_empty_intersection,
    )


def indexed_slices_intersection_indices_and_values(x1, x2):
    
    
    
    dim_0 = x1.dense_shape[0]
    x1_indices_expanded = tf.expand_dims(x1.indices, axis=1)
    x2_indices_expanded = tf.expand_dims(x2.indices, axis=1)
    x1_indices_count = x1_indices_expanded.shape[0]
    x2_indices_count = x2_indices_expanded.shape[0]
    x1_indices_one_hot = tf.scatter_nd(
        x1_indices_expanded,
        ones_bool((x1_indices_count,)),
        (dim_0,),
    )
    x2_indices_one_hot = tf.scatter_nd(
        x2_indices_expanded,
        ones_bool((x2_indices_count,)),
        (dim_0,),
    )
    intersection_indices = tf.squeeze(
        tf.where(tf.math.logical_and(x1_indices_one_hot, x2_indices_one_hot)),
        axis=-1,
    )
    intersection_indices_count = tf.shape(intersection_indices)[0]

    def empty_intersection():
        return (
            intersection_indices,
            empty_tensor((0,) + x1.values.shape[1:], x1.dtype),
            empty_tensor((0,) + x2.values.shape[1:], x2.dtype),
        )

    def non_empty_intersection():
        
        def values_for_intersection(indices_expanded, indices_count, values):
            indices_indices = tf.scatter_nd(
                indices_expanded,
                tf.range(indices_count),
                (dim_0,),
            )
            to_intersection_indices = tf.gather(
                indices_indices, intersection_indices
            )
            return tf.gather(values, to_intersection_indices)

        
        x1_values_for_intersection = tf.cond(
            tf.equal(x1_indices_count, intersection_indices_count),
            lambda: x1.values,
            lambda: values_for_intersection(
                x1_indices_expanded, x1_indices_count, x1.values
            ),
        )
        x2_values_for_intersection = tf.cond(
            tf.equal(x2_indices_count, intersection_indices_count),
            lambda: x2.values,
            lambda: values_for_intersection(
                x2_indices_expanded, x2_indices_count, x2.values
            ),
        )

        return (
            intersection_indices,
            x1_values_for_intersection,
            x2_values_for_intersection,
        )

    return tf.cond(
        tf.equal(intersection_indices_count, 0),
        empty_intersection,
        non_empty_intersection,
    )


def densifying_unary(default_value):
    

    def wrap_densifying_unary(func):
        @functools.wraps(func)
        def sparse_wrapper(x, *args, **kwargs):
            if isinstance(x, tf.SparseTensor):
                sparse_output = sparse_with_values(
                    x, func(x.values, *args, **kwargs)
                )
                return sparse_to_dense(
                    sparse_output,
                    tf.cast(default_value, sparse_output.values.dtype),
                )
            elif isinstance(x, tf.IndexedSlices):
                sparse_output_values = func(x.values, *args, **kwargs)
                output = tf.fill(
                    x.dense_shape,
                    tf.cast(default_value, sparse_output_values.dtype),
                )
                return tf.tensor_scatter_nd_update(
                    output, tf.expand_dims(x.indices, 1), sparse_output_values
                )
            return func(x, *args, **kwargs)

        return sparse_wrapper

    return wrap_densifying_unary


def elementwise_unary(func):
    

    @functools.wraps(func)
    def sparse_wrapper(x, *args, **kwargs):
        if isinstance(x, tf.SparseTensor):
            return sparse_with_values(x, func(x.values, *args, **kwargs))
        elif isinstance(x, tf.IndexedSlices):
            return tf.IndexedSlices(
                func(x.values, *args, **kwargs), x.indices, x.dense_shape
            )
        else:
            return func(x, *args, **kwargs)

    return sparse_wrapper


def elementwise_binary_union(sparse_op, densify_mixed=False):
    

    def wrap_elementwise_binary_union(func):
        @functools.wraps(func)
        def sparse_wrapper(x1, x2):
            if isinstance(x1, tf.SparseTensor):
                if isinstance(x2, tf.SparseTensor):
                    
                    if x1.indices is x2.indices:
                        return sparse_with_values(
                            x1, func(x1.values, x2.values)
                        )
                    else:
                        output = sparse_op(x1, x2)
                        output.set_shape(x1.shape)
                        return output
                else:
                    
                    if densify_mixed:
                        x1 = sparse_to_dense(x1)
                    else:
                        if not hasattr(x2, "shape") or len(x2.shape) == 0:
                            
                            x2 = broadcast_scalar_to_sparse_shape(x2, x1)
                        return sparse_op(x1, x2)
            elif isinstance(x2, tf.SparseTensor):
                
                if densify_mixed:
                    x2 = sparse_to_dense(x2)
                else:
                    if not hasattr(x1, "shape") or len(x1.shape) == 0:
                        
                        x1 = broadcast_scalar_to_sparse_shape(x1, x2)
                    return sparse_op(x1, x2)
            elif isinstance(x1, tf.IndexedSlices):
                if isinstance(x2, tf.IndexedSlices):
                    
                    if x1.indices is x2.indices:
                        return tf.IndexedSlices(
                            func(x1.values, x2.values),
                            x1.indices,
                            x1.dense_shape,
                        )
                    else:
                        
                        (
                            union_indices,
                            x1_values_for_union,
                            x2_values_for_union,
                        ) = indexed_slices_union_indices_and_values(
                            x1, x2.indices, x2.values
                        )
                        
                        return tf.IndexedSlices(
                            func(
                                x1_values_for_union,
                                x2_values_for_union,
                            ),
                            union_indices,
                            x1.dense_shape,
                        )
                else:
                    
                    x1 = tf.convert_to_tensor(x1)
            elif isinstance(x2, tf.IndexedSlices):
                
                x2 = tf.convert_to_tensor(x2)
            return func(x1, x2)

        return sparse_wrapper

    return wrap_elementwise_binary_union


def elementwise_binary_intersection(func):
    

    @functools.wraps(func)
    def sparse_wrapper(x1, x2):
        if isinstance(x1, tf.SparseTensor):
            if isinstance(x2, tf.SparseTensor):
                
                if x1.indices is x2.indices:
                    return sparse_with_values(x1, func(x1.values, x2.values))
                else:
                    
                    (
                        intersection_indices,
                        x1_values_for_intersection,
                        x2_values_for_intersection,
                    ) = sparse_intersection_indices_and_values(x1, x2)
                    
                    output = tf.SparseTensor(
                        intersection_indices,
                        func(
                            x1_values_for_intersection,
                            x2_values_for_intersection,
                        ),
                        x1.dense_shape,
                    )
                    output.set_shape(x1.shape)
                    return output
            else:
                
                if not hasattr(x2, "shape") or len(x2.shape) == 0:
                    
                    return sparse_with_values(x1, func(x1.values, x2))
                else:
                    
                    return sparse_with_values(
                        x1, func(x1.values, tf.gather_nd(x2, x1.indices))
                    )
        elif isinstance(x2, tf.SparseTensor):
            
            if not hasattr(x1, "shape") or len(x1.shape) == 0:
                
                return sparse_with_values(x2, func(x1, x2.values))
            else:
                
                return sparse_with_values(
                    x2, func(tf.gather_nd(x1, x2.indices), x2.values)
                )
        elif isinstance(x1, tf.IndexedSlices):
            if isinstance(x2, tf.IndexedSlices):
                
                if x1.indices is x2.indices:
                    return tf.IndexedSlices(
                        func(x1.values, x2.values), x1.indices, x1.dense_shape
                    )
                else:
                    
                    (
                        intersection_indices,
                        x1_values_for_intersection,
                        x2_values_for_intersection,
                    ) = indexed_slices_intersection_indices_and_values(x1, x2)
                    
                    return tf.IndexedSlices(
                        func(
                            x1_values_for_intersection,
                            x2_values_for_intersection,
                        ),
                        intersection_indices,
                        x1.dense_shape,
                    )
            else:
                
                if not hasattr(x2, "shape") or len(x2.shape) == 0:
                    
                    return tf.IndexedSlices(
                        func(x1.values, x2), x1.indices, x1.dense_shape
                    )
                else:
                    
                    return tf.IndexedSlices(
                        func(x1.values, tf.gather(x2, x1.indices)),
                        x1.indices,
                        x1.dense_shape,
                    )
        elif isinstance(x2, tf.IndexedSlices):
            
            if not hasattr(x1, "shape") or len(x1.shape) == 0:
                
                return tf.IndexedSlices(
                    func(x1, x2.values), x2.indices, x2.dense_shape
                )
            else:
                
                return tf.IndexedSlices(
                    func(tf.gather(x1, x2.indices), x2.values),
                    x2.indices,
                    x2.dense_shape,
                )
        
        return func(x1, x2)

    return sparse_wrapper


def elementwise_division(func):
    

    @functools.wraps(func)
    def sparse_wrapper(x1, x2):
        if isinstance(x1, tf.SparseTensor):
            if isinstance(x2, tf.SparseTensor):
                
                
                
                x1 = sparse_to_dense(x1)
                x2 = sparse_to_dense(x2)
            else:
                
                if not hasattr(x2, "shape") or len(x2.shape) == 0:
                    
                    return sparse_with_values(x1, func(x1.values, x2))
                else:
                    
                    x2_zeros_and_nans = tf.equal(x2, 0)
                    if not tf.as_dtype(x2.dtype).is_integer:
                        x2_zeros_and_nans = tf.math.logical_or(
                            x2_zeros_and_nans, tf.math.is_nan(x2)
                        )

                    def func_for_x1_indices():
                        
                        return sparse_with_values(
                            x1, func(x1.values, tf.gather_nd(x2, x1.indices))
                        )

                    def func_for_union_indices():
                        
                        x2_zeros_and_nan_indices = tf.where(x2_zeros_and_nans)
                        (
                            union_indices,
                            x1_values_for_union,
                            _,
                        ) = sparse_union_indices_and_values(
                            x1, x2_zeros_and_nan_indices
                        )
                        output = tf.SparseTensor(
                            union_indices,
                            func(
                                x1_values_for_union,
                                tf.gather_nd(x2, union_indices),
                            ),
                            x1.dense_shape,
                        )
                        output.set_shape(x1.shape)
                        return output

                    return tf.cond(
                        tf.reduce_any(x2_zeros_and_nans),
                        func_for_union_indices,
                        func_for_x1_indices,
                    )
        elif isinstance(x2, tf.SparseTensor):
            
            
            x2 = sparse_to_dense(x2)
        elif isinstance(x1, tf.IndexedSlices):
            if isinstance(x2, tf.IndexedSlices):
                
                
                
                x1 = tf.convert_to_tensor(x1)
                x2 = tf.convert_to_tensor(x2)
            else:
                
                if not hasattr(x2, "shape") or len(x2.shape) == 0:
                    
                    return tf.IndexedSlices(
                        func(x1.values, x2), x1.indices, x1.dense_shape
                    )
                else:
                    
                    x2_zeros_and_nans = tf.equal(x2, 0)
                    if not tf.as_dtype(x2.dtype).is_integer:
                        x2_zeros_and_nans = tf.math.logical_or(
                            x2_zeros_and_nans, tf.math.is_nan(x2)
                        )
                    x2_zeros_and_nans = tf.reduce_any(
                        x2_zeros_and_nans, axis=tuple(range(1, x2.shape.rank))
                    )

                    def func_for_x1_indices():
                        
                        return tf.IndexedSlices(
                            func(x1.values, tf.gather(x2, x1.indices)),
                            x1.indices,
                            x1.dense_shape,
                        )

                    def func_for_union_indices():
                        x2_zeros_and_nan_indices = tf.squeeze(
                            tf.where(x2_zeros_and_nans), axis=-1
                        )
                        
                        (
                            union_indices,
                            x1_values_for_union,
                            _,
                        ) = indexed_slices_union_indices_and_values(
                            x1, x2_zeros_and_nan_indices
                        )
                        return tf.IndexedSlices(
                            func(
                                x1_values_for_union,
                                tf.gather(x2, union_indices),
                            ),
                            union_indices,
                            x1.dense_shape,
                        )

                    return tf.cond(
                        tf.reduce_any(x2_zeros_and_nans),
                        func_for_union_indices,
                        func_for_x1_indices,
                    )
        elif isinstance(x2, tf.IndexedSlices):
            
            
            x2 = tf.convert_to_tensor(x2)
        
        return func(x1, x2)

    return sparse_wrapper