code: 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,
      precision: tuple[PrecisionType, PrecisionType] | None,
      preferred_element_type: DType | None,
      _in_avals: Sequence[core.ShapedArray],
      _out_aval: core.ShapedArray,
  ):
      """Implementation of lax.conv_general_dilated_p using XlaConv."""
      # In presence of shape polymorphism, lhs.shape and rhs.shape may contain
      # None. The actual dimension polynomial shapes are in _in_avals.
      del precision  # Unused arguments.
      lhs_shape, rhs_shape = _in_avals[0].shape, _in_avals[1].shape
      out_shape = _out_aval.shape
      _validate_spatial_dimensions(lhs, lhs_shape, rhs, rhs_shape)
      is_conv1d = len(lhs_shape) - 2 == 1
  
      tf_window_strides = _normalize_window_strides(window_strides)
      (
          padding,
          lhs_dilation,
          rhs_dilation,
      ) = _normalize_padding_and_dilations(
          padding, lhs_dilation, rhs_dilation, is_conv1d
      )
  
      lhs, lhs_shape, rhs, rhs_shape = _transpose_for_tf_conv(
          lhs, lhs_shape, rhs, rhs_shape, dimension_numbers
      )
      in_channels = lhs_shape[-1]
      *rhs_spatial_shapes, _, rhs_out_channel = rhs_shape
  
      is_transpose = any(d != 1 for d in lhs_dilation)
      is_atrous = any(d != 1 for d in rhs_dilation)
      is_depthwise = (
          in_channels == feature_group_count and feature_group_count > 1
      )
      _validate_conv_features(
          is_transpose,
          is_atrous,
          is_depthwise,
          feature_group_count,
          batch_group_count,
          preferred_element_type,
          lhs.dtype.as_numpy_dtype,
      )
  
      rhs_dilated_shape = [
          (k - 1) * r + 1
          for k, r in zip(rhs_spatial_shapes, rhs_dilation)
      ]
      output_perm = dimension_numbers[2]
  
      if is_transpose:
          padding_type = _conv_transpose_pads_to_padtype(
              rhs_spatial_shapes, lhs_dilation, padding
          )
      else:
          padding_type = pads_to_padtype(
              lhs_shape[1:3], rhs_dilated_shape, window_strides, padding
          )
          # We only manually pad if we aren't using a transposed convolutions.
          if padding_type == "EXPLICIT":
              lhs, lhs_shape, padding = _check_pad_spatial_dims(
                  lhs, lhs_shape, padding
              )
              padding_type = padding
  
      if padding_type != "SAME" and any(
          l < r for l, r in zip(lhs_shape[1:3], rhs_dilated_shape)
      ):
          # If the input shape is smaller than the filter shape in a spatial dimension,
          # lax returns only zeros while tf.conv2d returns an error.
          # We thus return zeros to make sure the behavior is consistent.
          return tf.broadcast_to(
              tf.constant(0, dtype=tf.float32),
              jax2tf._eval_shape(out_shape),
          )
  
      if is_depthwise:
          # Reshape filter from
          # [filter_height, filter_width, 1, in_channels * channel_multiplier] to
          # [filter_height, filter_width, in_channels, channel_multiplier].
          new_rhs_shape = tuple(rhs_spatial_shapes) + (
              in_channels,
              rhs_out_channel // in_channels,
          )
          output = tf.nn.depthwise_conv2d(
              input=lhs,
              filter=tf.reshape(rhs, jax2tf._eval_shape(new_rhs_shape)),
              strides=tf_window_strides,
              padding=padding_type,
              dilations=rhs_dilation,
          )
  
      elif is_transpose:
          # tf.nn.conv2d_transpose requires a transposed filter.
          rhs_t = tf.reverse(rhs, [0, 1])
          rhs_t = tf.transpose(rhs_t, (0, 1, 3, 2))
  
          # We should transpose `out_shape` to "NHWC", which is what TF expects.
          # First transpose to "NCHW".
          if is_conv1d:
              tf_out_shape = tuple(
                  out_shape[i] for i in output_perm
              ) + (1,)
          else:
              tf_out_shape = tuple(out_shape[i] for i in output_perm)
          # Then transpose "NCHW" to "NHWC".
          tf_out_shape = tuple(tf_out_shape[i] for i in (0, 2, 3, 1))
          output = tf.nn.conv2d_transpose(
              input=lhs,
              filters=rhs_t,
              output_shape=jax2tf._eval_shape(tf_out_shape),
              strides=lhs_dilation,
              padding=padding_type,
          )
  
      else:
          output = tf.nn.conv2d(
              input=lhs,
              filters=rhs,
              strides=tf_window_strides,
              padding=padding_type,
              dilations=rhs_dilation,
          )
  
      # TF outputs in format "NHWC", so convert to "NCHW", which is lax's default
      # format.
      output = tf.transpose(output, (0, 3, 1, 2))  # "NHWC" --> "NCHW"
      if is_conv1d:
          output = output[:, :, :, 0]
      # To determine the right permutation, we compute the inverse permutation of
      # `output_perm`, so that when `output_perm` is applied to `output`, we obtain
      # the outpt in NCHW format.
      inverse_perm = _invert_permutation(output_perm)
      output = tf.transpose(
          output, inverse_perm
      )  # "NCHW" -> desired output shape.
      return output
  
imports: from __future__ import annotations
  import builtins
  from collections.abc import Sequence
  import dataclasses
  from functools import partial, wraps
  import math
  import string
  from typing import Any, Callable, Optional
  from jax._src import core
  from jax import lax
  from jax._src.lax import slicing
  from jax._src import dtypes
  from jax._src import util
  from jax.experimental.jax2tf import jax2tf
  import numpy
  import tensorflow
label_desc: Use image node_id 14 for calling a global function with example usage: _conv_general_dilated(lhs, rhs) and returns: output, tf