code: def _get_kernel_shape(
      self, input_shape, kernel_size, module, n_features
  ):
      if module == nn.Conv:
          kernel_shape = kernel_size + (input_shape[-1], n_features)
      elif module == nn.ConvLocal:
          kernel_shape = input_shape[1:-1] + (
              input_shape[-1] * np.prod(kernel_size),
              n_features,
          )
      else:
          raise ValueError(module)
      return kernel_shape
  
imports: import functools
  import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest, parameterized
  from jax import random
  from jax.nn import initializers
  from flax import linen
label_desc: Use image node_id 20 for calling the LinearTest obj's underlying member method code with example usage: obj._get_kernel_shape(input_shape, kernel_size, module, n_features) and returns: kernel_shape