code: def _get_net(
      W_std,
      b_std,
      filter_shape,
      is_conv,
      use_pooling,
      is_res,
      padding,
      phi,
      strides,
      width,
      is_ntk,
      proj_into_2d,
      pool_type,
      layer_norm,
      parameterization,
      s,
      use_dropout,
  ):
      if is_conv:
          # Select a random filter order.
          default_filter_spec = "HW"
          filter_specs = [
              "".join(p) for p in itertools.permutations("HWIO")
          ]
          filter_spec = prandom.choice(filter_specs)
          filter_shape = tuple(
              filter_shape[default_filter_spec.index(c)]
              for c in filter_spec
              if c in default_filter_spec
          )
          strides = tuple(
              strides[default_filter_spec.index(c)]
              for c in filter_spec
              if c in default_filter_spec
          )
  
          # Select the activation order.
          default_spec = "NHWC"
          if default_backend() == "tpu":
              # Keep batch dimension leading for TPU for batching to work.
              specs = [
                  "N" + "".join(p)
                  for p in itertools.permutations("CHW")
              ]
          else:
              specs = [
                  "".join(p) for p in itertools.permutations("NCHW")
              ]
          spec = prandom.choice(specs)
          input_shape = tuple(
              INPUT_SHAPE[default_spec.index(c)] for c in spec
          )
  
      else:
          input_shape = (INPUT_SHAPE[0], np.prod(INPUT_SHAPE[1:]))
          if default_backend() == "tpu":
              spec = "NC"
          else:
              spec = prandom.choice(["NC", "CN"])
              if spec.index("N") == 1:
                  input_shape = input_shape[::-1]
  
          filter_spec = None
  
      dimension_numbers = (spec, filter_spec, spec)
      batch_axis, channel_axis = spec.index("N"), spec.index("C")
  
      spec_fc = "".join(c for c in spec if c in ("N", "C"))
      batch_axis_fc, channel_axis_fc = spec_fc.index(
          "N"
      ), spec_fc.index("C")
  
      if not is_conv:
          batch_axis = batch_axis_fc
          channel_axis = channel_axis_fc
  
      if layer_norm:
          layer_norm = tuple(spec.index(c) for c in layer_norm)
  
      def fc(out_dim, s):
          return stax.Dense(
              out_dim=out_dim,
              W_std=W_std,
              b_std=b_std,
              parameterization=parameterization,
              s=s,
              batch_axis=batch_axis_fc,
              channel_axis=channel_axis_fc,
          )
  
      def conv(out_chan, s):
          return stax.Conv(
              out_chan=out_chan,
              filter_shape=filter_shape,
              strides=strides,
              padding=padding,
              W_std=W_std,
              b_std=b_std,
              dimension_numbers=dimension_numbers,
              parameterization=parameterization,
              s=s,
          )
  
      affine = conv(width, (s, s)) if is_conv else fc(width, (s, s))
      affine_bottom = (
          conv(width, (1, s)) if is_conv else fc(width, (1, s))
      )
  
      rate = np.random.uniform(0.5, 0.9)
      dropout = stax.Dropout(rate, mode="train")
  
      if pool_type == "AVG":
          pool_fn = stax.AvgPool
          global_pool_fn = stax.GlobalAvgPool
      elif pool_type == "SUM":
          pool_fn = stax.SumPool
          global_pool_fn = stax.GlobalSumPool
      else:
          raise ValueError(pool_type)
  
      if use_pooling:
          pool_or_identity = pool_fn(
              (2, 3),
              None,
              "SAME" if padding == "SAME" else "CIRCULAR",
              batch_axis=batch_axis,
              channel_axis=channel_axis,
          )
      else:
          pool_or_identity = stax.Identity()
      dropout_or_identity = dropout if use_dropout else stax.Identity()
      layer_norm_or_identity = (
          stax.Identity()
          if layer_norm is None
          else stax.LayerNorm(
              axis=layer_norm,
              batch_axis=batch_axis,
              channel_axis=channel_axis,
          )
      )
      res_unit = stax.serial(
          dropout_or_identity, affine, pool_or_identity
      )
      if is_res:
          block = stax.serial(
              affine_bottom,
              stax.FanOut(2),
              stax.parallel(stax.Identity(), res_unit),
              stax.FanInSum(),
              layer_norm_or_identity,
              phi,
          )
      else:
          block = stax.serial(
              affine_bottom, res_unit, layer_norm_or_identity, phi
          )
  
      if proj_into_2d == "FLAT":
          proj_layer = stax.Flatten(batch_axis, batch_axis_fc)
      elif proj_into_2d == "POOL":
          proj_layer = global_pool_fn(batch_axis, channel_axis)
      elif proj_into_2d.startswith("ATTN"):
          n_heads = int(jnp.sqrt(width))
          n_chan_val = int(jnp.round(float(width) / n_heads))
          proj_layer = stax.serial(
              stax.GlobalSelfAttention(
                  n_chan_out=width,
                  n_chan_key=width,
                  n_chan_val=n_chan_val,
                  n_heads=n_heads,
                  linear_scaling=True,
                  W_key_std=W_std,
                  W_value_std=W_std,
                  W_query_std=W_std,
                  W_out_std=1.0,
                  b_std=b_std,
                  batch_axis=batch_axis,
                  channel_axis=channel_axis,
              ),
              stax.Flatten(batch_axis, batch_axis_fc),
          )
      else:
          raise ValueError(proj_into_2d)
  
      readout = stax.serial(
          proj_layer,
          fc(1 if is_ntk else width, (s, 1 if is_ntk else s)),
      )
  
      device_count = -1 if spec.index("N") == 0 else 0
  
      net = stax.serial(block, readout)
      return net, input_shape, device_count, channel_axis_fc
  
imports: import functools
  import itertools
  import random
  from absl.testing import absltest
  from jax import config
  from jax import default_backend
  from jax import jit
  from jax import random
  from jax.example_libraries import stax
  import jax.numpy
  import neural_tangents
  from neural_tangents import stax
  from neural_tangents._src.empirical import _DEFAULT_TESTING_NTK_IMPLEMENTATION
  from tests import test_utils
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: _get_net(W_std, b_std, filter_shape, is_conv, use_pooling, is_res, padding, phi, strides, width, is_ntk, proj_into_2d, pool_type, layer_norm, parameterization, s, use_dropout) and returns: net,  input_shape,  device_count,  channel_axis_fc, stax, stax