code: def get_norm_module(
      module: nn.Module,
      causal: bool = False,
      norm: str = "none",
      **norm_kwargs,
  ):
      """Return the proper normalization module. If causal is True, this will ensure the returned
      module is causal, or return an error if the normalization doesn't support causal evaluation.
      """
      assert norm in CONV_NORMALIZATIONS
      if norm == "time_group_norm":
          if causal:
              raise ValueError(
                  "GroupNorm doesn't support causal evaluation."
              )
          assert isinstance(module, nn.modules.conv._ConvNd)
          return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
      else:
          return nn.Identity()
  
imports: import math
  import typing
  import warnings
  import torch
  from torch import nn
  from torch.nn import functional
  from torch.nn.utils import spectral_norm, weight_norm
label_desc: Use image node_id 2 for calling a global function with example usage: get_norm_module(module, causal, norm) and returns: nn, nn