code: def make_downsampling_layer(
      strategy: Union[str, DownsamplingStrategy],
      output_channels: int,
  ) -> hk.SupportsCall:
      """Returns a sequence of modules corresponding to the desired downsampling."""
      strategy = DownsamplingStrategy(strategy)
  
      if strategy is DownsamplingStrategy.AVG_POOL:
          return hk.AvgPool(
              window_shape=(3, 3, 1), strides=(2, 2, 1), padding="SAME"
          )
  
      elif strategy is DownsamplingStrategy.CONV:
          return hk.Sequential(
              [
                  hk.Conv2D(
                      output_channels,
                      kernel_shape=3,
                      stride=2,
                      w_init=hk.initializers.TruncatedNormal(1e-2),
                  ),
              ]
          )
  
      elif strategy is DownsamplingStrategy.LAYERNORM_RELU_CONV:
          return hk.Sequential(
              [
                  hk.LayerNorm(
                      axis=(-3, -2, -1),
                      create_scale=True,
                      create_offset=True,
                      eps=1e-6,
                  ),
                  jax.nn.relu,
                  hk.Conv2D(
                      output_channels,
                      kernel_shape=3,
                      stride=2,
                      w_init=hk.initializers.TruncatedNormal(1e-2),
                  ),
              ]
          )
  
      elif strategy is DownsamplingStrategy.CONV_MAX:
          return hk.Sequential(
              [
                  hk.Conv2D(output_channels, kernel_shape=3, stride=1),
                  hk.MaxPool(
                      window_shape=(3, 3, 1),
                      strides=(2, 2, 1),
                      padding="SAME",
                  ),
              ]
          )
      else:
          raise ValueError(
              "Unrecognized downsampling strategy. Expected one of"
              f" {[strategy.value for strategy in DownsamplingStrategy]}"
              f" but received {strategy}."
          )
  
file_path: acme/acme/jax/networks/resnet.py
imports: import enum
  import functools
  from typing import Callable, Sequence, Union
  import haiku
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: make_downsampling_layer(strategy, output_channels) and returns: hk, hk, hk, hk