def _conv_general_dilated(lhs, rhs, *, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers: lax.ConvDimensionNumbers, feature_group_count: int, batch_group_count: int, lhs_shape: Sequence[int], rhs_shape: Sequence[int], precision: Optional[Tuple[(PrecisionType, PrecisionType)]], preferred_element_type: Optional[DType], _in_avals: Sequence[core.AbstractValue], _out_aval: core.AbstractValue):
    'Implementation of lax.conv_general_dilated_p using XlaConv.'
    out_tf_shape = _aval_to_tf_shape(_out_aval)
    if (not _enable_xla):
        return _try_tf_conv(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, batch_group_count, preferred_element_type, out_tf_shape)
    dnums_proto = _conv_general_dimension_numbers_proto(dimension_numbers)
    precision_config_proto = _precision_config_proto(precision)
    assert (batch_group_count == 1)

    def gen_conv(lhs, rhs, preferred_element_type: Optional[DType]):
        out = tfxla.conv(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dnums_proto, feature_group_count=feature_group_count, precision_config=precision_config_proto, preferred_element_type=preferred_element_type)
        out.set_shape(out_tf_shape)
        return out
    if np.issubdtype(_in_avals[0].dtype, np.complexfloating):
        if (preferred_element_type is not None):
            assert np.issubdtype(preferred_element_type, np.complexfloating)
            preferred_float_et = (np.float64 if (preferred_element_type == np.complex128) else np.float32)
        else:
            preferred_float_et = None
        (lhs_real, lhs_imag) = (tf.math.real(lhs), tf.math.imag(lhs))
        (rhs_real, rhs_imag) = (tf.math.real(rhs), tf.math.imag(rhs))
        k1 = gen_conv(_add(lhs_real, lhs_imag), rhs_real, preferred_float_et)
        k2 = gen_conv(lhs_real, tf.math.subtract(rhs_imag, rhs_real), preferred_float_et)
        k3 = gen_conv(lhs_imag, _add(rhs_real, rhs_imag), preferred_float_et)
        return tf.complex(tf.math.subtract(k1, k3), _add(k1, k2))
    else:
        return gen_conv(lhs, rhs, preferred_element_type)
Implementation of lax.conv_general_dilated_p using XlaConv.
def _conv_general_dilated(lhs, rhs, *, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers: lax.ConvDimensionNumbers, feature_group_count: int, batch_group_count: int, lhs_shape: Sequence[int], rhs_shape: Sequence[int], precision: Optional[Tuple[(PrecisionType, PrecisionType)]], preferred_element_type: Optional[DType], _in_avals: Sequence[core.AbstractValue], _out_aval: core.AbstractValue):
    
    out_tf_shape = _aval_to_tf_shape(_out_aval)
    if (not _enable_xla):
        return _try_tf_conv(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, batch_group_count, preferred_element_type, out_tf_shape)
    dnums_proto = _conv_general_dimension_numbers_proto(dimension_numbers)
    precision_config_proto = _precision_config_proto(precision)
    assert (batch_group_count == 1)

    def gen_conv(lhs, rhs, preferred_element_type: Optional[DType]):
        out = tfxla.conv(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dnums_proto, feature_group_count=feature_group_count, precision_config=precision_config_proto, preferred_element_type=preferred_element_type)
        out.set_shape(out_tf_shape)
        return out
    if np.issubdtype(_in_avals[0].dtype, np.complexfloating):
        if (preferred_element_type is not None):
            assert np.issubdtype(preferred_element_type, np.complexfloating)
            preferred_float_et = (np.float64 if (preferred_element_type == np.complex128) else np.float32)
        else:
            preferred_float_et = None
        (lhs_real, lhs_imag) = (tf.math.real(lhs), tf.math.imag(lhs))
        (rhs_real, rhs_imag) = (tf.math.real(rhs), tf.math.imag(rhs))
        k1 = gen_conv(_add(lhs_real, lhs_imag), rhs_real, preferred_float_et)
        k2 = gen_conv(lhs_real, tf.math.subtract(rhs_imag, rhs_real), preferred_float_et)
        k3 = gen_conv(lhs_imag, _add(rhs_real, rhs_imag), preferred_float_et)
        return tf.complex(tf.math.subtract(k1, k3), _add(k1, k2))
    else:
        return gen_conv(lhs, rhs, preferred_element_type)