def conv2d(x, filters, strides, padding, data_format='NHWC', dilations=1):
    strides = (([strides] * 2) if isinstance(strides, int) else strides)
    dilations = (([dilations] * 2) if isinstance(dilations, int) else dilations)
    return jlax.conv_general_dilated(x, filters, strides, padding, None, dilations, (data_format, 'HWIO', data_format))
def conv2d(x, filters, strides, padding, data_format='NHWC', dilations=1):
    strides = (([strides] * 2) if isinstance(strides, int) else strides)
    dilations = (([dilations] * 2) if isinstance(dilations, int) else dilations)
    return jlax.conv_general_dilated(x, filters, strides, padding, None, dilations, (data_format, 'HWIO', data_format))