class_docstr: General N-dimensional convolutional.
class_docstr_tok: ["General","N-dimensional","convolutional","."]
docstr: Connects ``ConvND`` layer.
  
  Args:
    inputs: An array of shape ``[spatial_dims, C]`` and rank-N+1 if unbatched,
      or an array of shape ``[N, spatial_dims, C]`` and rank-N+2 if batched.
    precision: Optional :class:`jax.lax.Precision` to pass to
      :func:`jax.lax.conv_general_dilated`.
  
  Returns:
    An array of shape ``[spatial_dims, output_channels]`` and rank-N+1 if
      unbatched, or an array of shape ``[N, spatial_dims, output_channels]``
      and rank-N+2 if batched.
docstr_tok: ["Connects","``","ConvND","``","layer",".","Args",":","inputs",":","An","array","of","shape","``","[","spatial_dims",",","C","]","``","and","rank-N+1","if","unbatched",",","or","an","array","of","shape","``","[","N",",","spatial_dims",",","C","]","``","and","rank-N+2","if","batched",".","precision",":","Optional",":","class",":","`","jax.lax.Precision","`","to","pass","to",":","func",":","`","jax.lax.conv_general_dilated","`",".","Returns",":","An","array","of","shape","``","[","spatial_dims",",","output_channels","]","``","and","rank-N+1","if","unbatched",",","or","an","array","of","shape","``","[","N",",","spatial_dims",",","output_channels","]","``","and","rank-N+2","if","batched","."]
code: def __call__(
      self,
      inputs: jax.Array,
      *,
      precision: Optional[lax.Precision] = None,
  ) -> jax.Array:
      """Connects ``ConvND`` layer.
  
      Args:
        inputs: An array of shape ``[spatial_dims, C]`` and rank-N+1 if unbatched,
          or an array of shape ``[N, spatial_dims, C]`` and rank-N+2 if batched.
        precision: Optional :class:`jax.lax.Precision` to pass to
          :func:`jax.lax.conv_general_dilated`.
  
      Returns:
        An array of shape ``[spatial_dims, output_channels]`` and rank-N+1 if
          unbatched, or an array of shape ``[N, spatial_dims, output_channels]``
          and rank-N+2 if batched.
      """
      unbatched_rank = self.num_spatial_dims + 1
      allowed_ranks = [unbatched_rank, unbatched_rank + 1]
      if inputs.ndim not in allowed_ranks:
          raise ValueError(
              f"Input to ConvND needs to have rank in {allowed_ranks},"
              f" but input has shape {inputs.shape}."
          )
  
      unbatched = inputs.ndim == unbatched_rank
      if unbatched:
          inputs = jnp.expand_dims(inputs, axis=0)
  
      if (
          inputs.shape[self.channel_index] % self.feature_group_count
          != 0
      ):
          raise ValueError(
              f"Inputs channels {inputs.shape[self.channel_index]} "
              f"should be a multiple of feature_group_count "
              f"{self.feature_group_count}"
          )
      w_shape = self.kernel_shape + (
          inputs.shape[self.channel_index] // self.feature_group_count,
          self.output_channels,
      )
  
      if self.mask is not None and self.mask.shape != w_shape:
          raise ValueError(
              "Mask needs to have the same shape as weights. "
              f"Shapes are: {self.mask.shape}, {w_shape}"
          )
  
      w_init = self.w_init
      if w_init is None:
          fan_in_shape = np.prod(w_shape[:-1])
          stddev = 1.0 / np.sqrt(fan_in_shape)
          w_init = hk.initializers.TruncatedNormal(stddev=stddev)
      w = hk.get_parameter("w", w_shape, inputs.dtype, init=w_init)
  
      if self.mask is not None:
          w *= self.mask
  
      out = lax.conv_general_dilated(
          inputs,
          w,
          window_strides=self.stride,
          padding=self.padding,
          lhs_dilation=self.lhs_dilation,
          rhs_dilation=self.kernel_dilation,
          dimension_numbers=self.dimension_numbers,
          feature_group_count=self.feature_group_count,
          precision=precision,
      )
  
      if self.with_bias:
          if self.channel_index == -1:
              bias_shape = (self.output_channels,)
          else:
              bias_shape = (self.output_channels,) + (
                  1,
              ) * self.num_spatial_dims
          b = hk.get_parameter(
              "b", bias_shape, inputs.dtype, init=self.b_init
          )
          b = jnp.broadcast_to(b, out.shape)
          out = out + b
  
      if unbatched:
          out = jnp.squeeze(out, axis=0)
      return out
  
imports: from collections.abc import Sequence
  from typing import Optional, Union
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import pad
  from haiku._src import utils
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the ConvND obj's underlying member method code with example usage: obj.__call__(inputs) and returns: out