class_docstr: Convolution Module wrapping `lax.conv_general_dilated[_local]`.
  
  Attributes:
    features: number of convolution filters.
    kernel_size: shape of the convolutional kernel. For 1D convolution,
      the kernel size can be passed as an integer, which will be interpreted
      as a tuple of the single integer. For all other cases, it must be a
      sequence of integers.
    strides: an integer or a sequence of `n` integers, representing the
      inter-window strides (default: 1).
    padding: either the string `'SAME'`, the string `'VALID'`, the string
      `'CIRCULAR'` (periodic boundary conditions), or a sequence of `n` `(low,
      high)` integer pairs that give the padding to apply before and after each
      spatial dimension. A single int is interpeted as applying the same padding
      in all dims and passign a single int in a sequence causes the same padding
      to be used on both sides. `'CAUSAL'` padding for a 1D convolution will
      left-pad the convolution axis, resulting in same-sized output.
    input_dilation: an integer or a sequence of `n` integers, giving the
      dilation factor to apply in each spatial dimension of `inputs`
      (default: 1). Convolution with input dilation `d` is equivalent to
      transposed convolution with stride `d`.
    kernel_dilation: an integer or a sequence of `n` integers, giving the
      dilation factor to apply in each spatial dimension of the convolution
      kernel (default: 1). Convolution with kernel dilation
      is also known as 'atrous convolution'.
    feature_group_count: integer, default 1. If specified divides the input
      features into groups.
    use_bias: whether to add a bias to the output (default: True).
    mask: Optional mask for the weights during masked convolution. The mask must
          be the same shape as the convolution weight matrix.
    dtype: the dtype of the computation (default: infer from input and params).
    param_dtype: the dtype passed to parameter initializers (default: float32).
    precision: numerical precision of the computation see `jax.lax.Precision`
      for details.
    kernel_init: initializer for the convolutional kernel.
    bias_init: initializer for the bias.
docstr: Applies a (potentially unshared) convolution to the inputs.
  
  Args:
    inputs: input data with dimensions (*batch_dims, spatial_dims...,
      features). This is the channels-last convention, i.e. NHWC for a 2d
      convolution and NDHWC for a 3D convolution. Note: this is different from
      the input convention used by `lax.conv_general_dilated`, which puts the
      spatial dimensions last.
      Note: If the input has more than 1 batch dimension, all batch dimensions
      are flattened into a single dimension for the convolution and restored
      before returning.  In some cases directly vmap'ing the layer may yield
      better performance than this default flattening approach.  If the input
      lacks a batch dimension it will be added for the convolution and removed
      n return, an allowance made to enable writing single-example code.
  
  Returns:
    The convolved data.
docstr_tok: ["Applies","a","(","potentially","unshared",")","convolution","to","the","inputs",".","Args",":","inputs",":","input","data","with","dimensions","(","*","batch_dims",",","spatial_dims","...",",","features",")",".","This","is","the","channels-last","convention",",","i.e",".","NHWC","for","a","2d","convolution","and","NDHWC","for","a","3D","convolution",".","Note",":","this","is","different","from","the","input","convention","used","by","`","lax.conv_general_dilated","`",",","which","puts","the","spatial","dimensions","last",".","Note",":","If","the","input","has","more","than","1","batch","dimension",",","all","batch","dimensions","are","flattened","into","a","single","dimension","for","the","convolution","and","restored","before","returning",".","In","some","cases","directly","vmap'ing","the","layer","may","yield","better","performance","than","this","default","flattening","approach",".","If","the","input","lacks","a","batch","dimension","it","will","be","added","for","the","convolution","and","removed","n","return",",","an","allowance","made","to","enable","writing","single-example","code",".","Returns",":","The","convolved","data","."]
code: def __call__(self, inputs: Array) -> Array:
      """Applies a (potentially unshared) convolution to the inputs.
  
      Args:
        inputs: input data with dimensions (*batch_dims, spatial_dims...,
          features). This is the channels-last convention, i.e. NHWC for a 2d
          convolution and NDHWC for a 3D convolution. Note: this is different from
          the input convention used by `lax.conv_general_dilated`, which puts the
          spatial dimensions last.
          Note: If the input has more than 1 batch dimension, all batch dimensions
          are flattened into a single dimension for the convolution and restored
          before returning.  In some cases directly vmap'ing the layer may yield
          better performance than this default flattening approach.  If the input
          lacks a batch dimension it will be added for the convolution and removed
          n return, an allowance made to enable writing single-example code.
  
      Returns:
        The convolved data.
      """
  
      assert isinstance(self.kernel_size, tuple)
      kernel_size = self.kernel_size
  
      def maybe_broadcast(
          x: tp.Optional[tp.Union[int, tp.Sequence[int]]],
      ) -> tp.Tuple[int, ...]:
          if x is None:
              # backward compatibility with using None as sentinel for
              # broadcast 1
              x = 1
          if isinstance(x, int):
              return (x,) * len(kernel_size)
          return tuple(x)
  
      # Combine all input batch dimensions into a single leading batch axis.
      num_batch_dimensions = inputs.ndim - (len(kernel_size) + 1)
      if num_batch_dimensions != 1:
          input_batch_shape = inputs.shape[:num_batch_dimensions]
          total_batch_size = int(np.prod(input_batch_shape))
          flat_input_shape = (total_batch_size,) + inputs.shape[
              num_batch_dimensions:
          ]
          inputs = jnp.reshape(inputs, flat_input_shape)
  
      # self.strides or (1,) * (inputs.ndim - 2)
      strides = maybe_broadcast(self.strides)
      input_dilation = maybe_broadcast(self.input_dilation)
      kernel_dilation = maybe_broadcast(self.kernel_dilation)
  
      padding_lax = canonicalize_padding(self.padding, len(kernel_size))
      if padding_lax == "CIRCULAR":
          kernel_size_dilated = [
              (k - 1) * d + 1
              for k, d in zip(kernel_size, kernel_dilation)
          ]
          zero_pad: tp.List[tp.Tuple[int, int]] = [(0, 0)]
          pads = (
              zero_pad
              + [((k - 1) // 2, k // 2) for k in kernel_size_dilated]
              + [(0, 0)]
          )
          inputs = jnp.pad(inputs, pads, mode="wrap")
          padding_lax = "VALID"
      elif padding_lax == "CAUSAL":
          if len(kernel_size) != 1:
              raise ValueError(
                  "Causal padding is only implemented for 1D convolutions."
              )
          left_pad = kernel_dilation[0] * (kernel_size[0] - 1)
          pads = [(0, 0), (left_pad, 0), (0, 0)]
          inputs = jnp.pad(inputs, pads)
          padding_lax = "VALID"
  
      dimension_numbers = _conv_dimension_numbers(inputs.shape)
  
      # One shared convolutional kernel for all pixels in the output.
      assert self.in_features % self.feature_group_count == 0
  
      kernel = self.kernel
  
      if self.mask_fn is not None:
          kernel = self.mask_fn(kernel)
  
      bias = self.bias
  
      inputs, kernel, bias = dtypes.promote_dtype(
          inputs, kernel, bias, dtype=self.dtype
      )
  
      y = self.conv_general_dilated(
          inputs,
          kernel,
          strides,
          padding_lax,
          lhs_dilation=input_dilation,
          rhs_dilation=kernel_dilation,
          dimension_numbers=dimension_numbers,
          feature_group_count=self.feature_group_count,
          precision=self.precision,
      )
  
      if self.use_bias:
          bias = bias.reshape((1,) * (y.ndim - bias.ndim) + bias.shape)
          y += bias
  
      if num_batch_dimensions != 1:
          output_shape = input_batch_shape + y.shape[1:]
          y = jnp.reshape(y, output_shape)
      return y
  
imports: from __future__ import annotations
  import typing
  from types import MappingProxyType
  import jax
  import jax.numpy
  import numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import rnglib, variables
  from flax.experimental.nnx.nnx.module import Module
  from flax.experimental.nnx.nnx.nn import dtypes, initializers
label_desc: Use image node_id 2 for calling the Conv obj's underlying member method code with example usage: obj.__call__(inputs) and returns: y, tuple, unknown