def _select_and_gather_add(tangents: TfVal, operand: TfVal, select_prim: core.Primitive, window_dimensions: Sequence[int], window_strides: Sequence[int], base_dilation: Sequence[int], window_dilation: Sequence[int], padding: Sequence[Tuple[(int, int)]], _in_avals: Sequence[core.AbstractValue], _out_aval: core.AbstractValue):
    dtype = operand.dtype
    nbits = dtypes.finfo(dtype.as_numpy_dtype).bits
    max_bits = 64
    assert (nbits <= max_bits)
    double_word_reduction = ((nbits * 2) <= max_bits)
    const = (lambda dtype, x: tf.constant(np.array(x), dtype))
    if double_word_reduction:
        word_dtype = lax._UINT_DTYPES[nbits]
        double_word_dtype = lax._UINT_DTYPES[(nbits * 2)]

        def pack(a, b):
            a = _bitcast_convert_type(a, word_dtype)
            b = _bitcast_convert_type(b, word_dtype)
            a = _convert_element_type(a, new_dtype=double_word_dtype)
            b = _convert_element_type(b, new_dtype=double_word_dtype)
            a = tf.bitwise.left_shift(a, const(double_word_dtype, nbits))
            return tf.bitwise.bitwise_or(a, b)

        def fst(t):
            assert (t.dtype == double_word_dtype)
            st = _shift_right_logical(t, const(double_word_dtype, nbits))
            return _bitcast_convert_type(_convert_element_type(st, new_dtype=word_dtype), dtype)

        def snd(t):
            return _bitcast_convert_type(_convert_element_type(t, new_dtype=word_dtype), dtype)
    else:
        raise NotImplementedError(f'TODO: need to pack {(nbits * 2)} bits but this platform can only go up to {max_bits} bits.')
    assert ((select_prim is lax.ge_p) or (select_prim is lax.le_p)), select_prim

    def reducer(x, y):
        which = tf_impl[select_prim]
        return tf_impl[lax.select_p](which(fst(x), fst(y)), x=x, y=y)
    init = ((- np.inf) if (select_prim is lax.ge_p) else np.inf)
    init_identity = (lambda x: pack(const(dtype, init), const(dtype, 0)))
    out = _specialized_reduce_window(reducer, init_identity, pack(operand, tangents), window_dimensions=window_dimensions, window_strides=window_strides, padding=padding, base_dilation=base_dilation, window_dilation=window_dilation, _in_avals=_in_avals, _out_aval=_out_aval)
    return snd(out)
def _select_and_gather_add(tangents: TfVal, operand: TfVal, select_prim: core.Primitive, window_dimensions: Sequence[int], window_strides: Sequence[int], base_dilation: Sequence[int], window_dilation: Sequence[int], padding: Sequence[Tuple[(int, int)]], _in_avals: Sequence[core.AbstractValue], _out_aval: core.AbstractValue):
    dtype = operand.dtype
    nbits = dtypes.finfo(dtype.as_numpy_dtype).bits
    max_bits = 64
    assert (nbits <= max_bits)
    double_word_reduction = ((nbits * 2) <= max_bits)
    const = (lambda dtype, x: tf.constant(np.array(x), dtype))
    if double_word_reduction:
        word_dtype = lax._UINT_DTYPES[nbits]
        double_word_dtype = lax._UINT_DTYPES[(nbits * 2)]

        def pack(a, b):
            a = _bitcast_convert_type(a, word_dtype)
            b = _bitcast_convert_type(b, word_dtype)
            a = _convert_element_type(a, new_dtype=double_word_dtype)
            b = _convert_element_type(b, new_dtype=double_word_dtype)
            a = tf.bitwise.left_shift(a, const(double_word_dtype, nbits))
            return tf.bitwise.bitwise_or(a, b)

        def fst(t):
            assert (t.dtype == double_word_dtype)
            st = _shift_right_logical(t, const(double_word_dtype, nbits))
            return _bitcast_convert_type(_convert_element_type(st, new_dtype=word_dtype), dtype)

        def snd(t):
            return _bitcast_convert_type(_convert_element_type(t, new_dtype=word_dtype), dtype)
    else:
        raise NotImplementedError(f'TODO: need to pack {(nbits * 2)} bits but this platform can only go up to {max_bits} bits.')
    assert ((select_prim is lax.ge_p) or (select_prim is lax.le_p)), select_prim

    def reducer(x, y):
        which = tf_impl[select_prim]
        return tf_impl[lax.select_p](which(fst(x), fst(y)), x=x, y=y)
    init = ((- np.inf) if (select_prim is lax.ge_p) else np.inf)
    init_identity = (lambda x: pack(const(dtype, init), const(dtype, 0)))
    out = _specialized_reduce_window(reducer, init_identity, pack(operand, tangents), window_dimensions=window_dimensions, window_strides=window_strides, padding=padding, base_dilation=base_dilation, window_dilation=window_dilation, _in_avals=_in_avals, _out_aval=_out_aval)
    return snd(out)