code: def _build_network(input_shape, network, out_logits):
      if len(input_shape) == 1:
          assert network == FLAT
          return stax.serial(
              stax.Dense(4096, W_std=1.2, b_std=0.05),
              stax.Erf(),
              stax.Dense(out_logits, W_std=1.2, b_std=0.05),
          )
      elif len(input_shape) == 3:
          if network == POOLING:
              return stax.serial(
                  stax.Conv(
                      CONVOLUTION_CHANNELS,
                      (3, 3),
                      W_std=2.0,
                      b_std=0.05,
                  ),
                  stax.GlobalAvgPool(),
                  stax.Dense(out_logits, W_std=2.0, b_std=0.05),
              )
          elif network == FLAT:
              return stax.serial(
                  stax.Conv(
                      CONVOLUTION_CHANNELS,
                      (3, 3),
                      W_std=2.0,
                      b_std=0.05,
                  ),
                  stax.Flatten(),
                  stax.Dense(out_logits, 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: import math
  from absl.testing import absltest
  from jax import grad
  from jax import jit
  from jax import random
  from jax import vmap
  from jax import config
  from jax.example_libraries import optimizers
  from jax.flatten_util import ravel_pytree
  import jax.numpy
  import jax.tree_util
  import neural_tangents
  from neural_tangents import predict, stax
  from neural_tangents._src.predict import _is_on_cpu
  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) and returns: stax, stax, stax