class_docstr: N-D Depthwise Convolution Module.
class_docstr_tok: ["N-D","Depthwise","Convolution","Module","."]
code: def __call__(self, inputs: jax.Array) -> jax.Array:
      channel_index = hk.get_channel_index(self.data_format)
      w_shape = self.kernel_shape + (
          1,
          self.channel_multiplier * inputs.shape[channel_index],
      )
  
      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)
  
      out = lax.conv_general_dilated(
          inputs,
          w,
          window_strides=self.stride,
          padding=self.padding,
          lhs_dilation=self.lhs_dilation,
          rhs_dilation=self.rhs_dilation,
          dimension_numbers=self.dn,
          feature_group_count=inputs.shape[channel_index],
      )
  
      if self.with_bias:
          if channel_index == -1:
              b_shape = (
                  self.channel_multiplier * inputs.shape[channel_index],
              )
          else:
              b_shape = (
                  self.channel_multiplier * inputs.shape[channel_index],
                  1,
                  1,
              )
          b = hk.get_parameter(
              "b", b_shape, inputs.dtype, init=self.b_init
          )
          b = jnp.broadcast_to(b, out.shape)
          out = out + b
  
      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 utils
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the DepthwiseConvND obj's underlying member method code with example usage: obj.__call__(inputs) and returns: out