def _depthwise_conv2d(inputs, kernel, strides, padding):
    """Computes a depthwise conv2d in Jax.
    Args:
        inputs: an NHWC tensor with N=1.
        kernel: a [H", W", 1, C] tensor.
        strides: a 2d tensor.
        padding: "SAME" or "VALID".
    Returns:
        The depthwise convolution of inputs with kernel, as [H, W, C].
    """
    return jax.lax.conv_general_dilated(
        inputs,
        kernel,
        strides,
        padding,
        feature_group_count=inputs.shape[-1],
        dimension_numbers=("NHWC", "HWIO", "NHWC"),
    )
Computes a depthwise conv2d in Jax.
Args:
inputs: an NHWC tensor with N=1.
kernel: a [H", W", 1, C] tensor.
strides: a 2d tensor.
padding: "SAME" or "VALID".
Returns:
The depthwise convolution of inputs with kernel, as [H, W, C].
serl_launcher/serl_launcher/vision/data_augmentations.py