code: def _build_network(input_shape, network, out_logits, use_dropout):
      dropout = (
          stax.Dropout(0.9, mode="train")
          if use_dropout
          else stax.Identity()
      )
      if len(input_shape) == 1:
          if network != FLAT:
              raise absltest.SkipTest(
                  "Not testing pooling on FCN inputs."
              )
  
          return stax.serial(
              stax.Dense(WIDTH, W_std=2.0, b_std=0.5),
              dropout,
              stax.Dense(out_logits, W_std=2.0, b_std=0.5),
          )
  
      elif len(input_shape) == 3:
          if network == POOLING:
              return stax.serial(
                  stax.Conv(
                      CONVOLUTION_CHANNELS,
                      (2, 2),
                      W_std=2.0,
                      b_std=0.05,
                  ),
                  stax.GlobalAvgPool(),
                  dropout,
                  stax.Dense(out_logits, W_std=2.0, b_std=0.5),
              )
          elif network == FLAT:
              return stax.serial(
                  stax.Conv(
                      CONVOLUTION_CHANNELS,
                      (2, 2),
                      W_std=2.0,
                      b_std=0.05,
                  ),
                  stax.Flatten(),
                  dropout,
                  stax.Dense(out_logits, W_std=2.0, b_std=0.5),
              )
          elif network == INTERMEDIATE_CONV:
              return stax.Conv(
                  CONVOLUTION_CHANNELS, (2, 2), W_std=2.0, b_std=0.05
              )
          else:
              raise ValueError(
                  "Unexpected network type found: {}".format(network)
              )
  
      else:
          raise ValueError("Expected flat or image test input.")
  
imports: from functools import partial
  from absl.testing import absltest
  from jax import jit
  from jax import config
  from jax import random
  import jax.numpy
  from jax.tree_util import tree_map
  import neural_tangents
  from neural_tangents import stax
  from neural_tangents._src import batching
  from neural_tangents._src.empirical import _DEFAULT_TESTING_NTK_IMPLEMENTATION
  from tests import test_utils
label_desc: Use image node_id 1 for calling a global function with example usage: _build_network(input_shape, network, out_logits, use_dropout) and returns: stax, stax, stax, stax