code: def _conv2d(
      input,  # pylint: disable=redefined-builtin
      filters,
      strides,
      padding,
      data_format="NHWC",
      dilations=None,
      name=None,
  ):
      """tf.nn.conv2d implementation."""
      del name
      if not JAX_MODE:
          raise NotImplementedError(
              "tf.nn.conv2d not implemented in NumPy."
          )
  
      if dilations is not None:
          raise ValueError("Dilations not yet supported")
  
      channel_index = 1 if data_format.startswith("NC") else 3
  
      window_strides = _get_sequence(
          strides, 2, channel_index, "strides"
      )
      if window_strides[0] != 1:
          raise ValueError(
              f"Stride != 1 not supported for batch dimension. `strides`: {strides} "
              f"`data_format`: {data_format}"
          )
      if window_strides[channel_index] != 1:
          raise ValueError(
              "Stride != 1 not supported for channel dimension. `strides`: "
              f"{strides} "
              f"`data_format`: {data_format}"
          )
      window_strides = [
          e
          for i, e in enumerate(window_strides)
          if i not in [0, channel_index]
      ]
  
      if isinstance(padding, (list, tuple)):
          if padding[0] != (0, 0):
              raise ValueError(
                  f"Padding not supported for batch dimension. `padding`: {padding} "
                  f"`data_format`: {data_format}"
              )
          if padding[channel_index] != (0, 0):
              raise ValueError(
                  "Padding not supported for channel dimension. `padding`: "
                  f"{padding} "
                  f"`data_format`: {data_format}"
              )
          padding = [
              s
              for i, s in enumerate(padding)
              if i not in [0, channel_index]
          ]
  
      return jax.lax.conv_general_dilated(
          lhs=input,
          rhs=filters,
          window_strides=window_strides,
          padding=padding,
          dimension_numbers=(data_format, "HWIO", data_format),
      )
  
file_path: probability/tensorflow_probability/python/internal/backend/numpy/nn.py
imports: import collections
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import numpy_array
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import l2_normalize
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import log_softmax
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import reduce_logsumexp
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import reduce_mean
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softmax
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softplus
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import squared_difference
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import top_k
  from tensorflow_probability.python.internal.backend.numpy.ops import stop_gradient
label_desc: Use image node_id 2 for calling a global function with example usage: _conv2d(input, filters, strides, padding, data_format, dilations, name) and returns: jax