code: def jax_conv(
      inp,
      fltr,
      window_strides,
      padding,
      dimension_numbers,
      filter_dilation=None,
  ):
      """A wrapper around `lax.conv_general_dilated`.
  
      It requires `dimension_numbers` and disallows `inp_dilation`.
  
      Args:
        inp: an (N+2)-D array. The input of the convolution.
        fltr: an (N+2)-D array. The filter (i.e. kernel) of the convolution.
        window_strides: the strides for moving the convolution window.
        padding: a string, either 'VALID' or 'SAME'. The padding algorithm.
        dimension_numbers: a tuple of three strings encoding the data format of
          input, filter and output. 'I' means input; 'O' means output; 'C' means
          channel; other characters such as 'W', 'H' and 'D' means spatial
          dimensions.
        filter_dilation: the dilation rates for the filter. Dilating the filter
          means adding "holes" to the filter.
  
      Returns:
        An (N+2)-D array. The convolution result.
      """
      return lax.conv_general_dilated(
          inp,
          fltr,
          window_strides,
          padding,
          lhs_dilation=None,
          rhs_dilation=filter_dilation,
          dimension_numbers=dimension_numbers,
      )
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import jax.scipy.special
  import numpy
  import tensorflow
  import tensorflow_datasets
  from trax.fastmath import numpy
  from trax.shapes import signature
label_desc: Use image node_id 1 for calling a global function with example usage: jax_conv(inp, fltr, window_strides, padding, dimension_numbers, filter_dilation) and returns: lax