code: def _get_net_pool(
      width,
      is_ntk,
      pool_type,
      padding,
      filter_shape,
      strides,
      normalize_edges,
  ):
      W_std, b_std = 2.0**0.5, 0.5**0.5
      phi = stax.Relu()
      parameterization = "ntk"
  
      fc = functools.partial(
          stax.Dense,
          W_std=W_std / width if pool_type == "SUM" else W_std,
          b_std=b_std,
          parameterization=parameterization,
      )
  
      conv = functools.partial(
          stax.Conv,
          filter_shape=filter_shape,
          strides=None,
          padding="SAME",
          W_std=W_std / np.prod(filter_shape)
          if pool_type == "SUM"
          else W_std,
          b_std=b_std,
          parameterization=parameterization,
      )
  
      if pool_type == "AVG":
          pool_fn = functools.partial(
              stax.AvgPool, normalize_edges=normalize_edges
          )
          global_pool_fn = stax.GlobalAvgPool
      elif pool_type == "SUM":
          pool_fn = stax.SumPool
          global_pool_fn = stax.GlobalSumPool
      else:
          raise ValueError(pool_type)
  
      pool = pool_fn(filter_shape, strides, padding)
  
      device_count = -1
  
      return (
          stax.serial(
              conv(width),
              phi,
              pool,
              conv(width),
              phi,
              global_pool_fn(),
              fc(1 if is_ntk else width),
          ),
          INPUT_SHAPE,
          device_count,
          -1,
      )
  
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 3 for calling a global function with example usage: _get_net_pool(width, is_ntk, pool_type, padding, filter_shape, strides, normalize_edges) and returns: stax,  INPUT_SHAPE,  device_count,  unknown