code: def test_fan_in_fc(
      self, same_inputs, axis, n_branches, get, branch_in, fan_in_mode
  ):
      if fan_in_mode in ["FanInSum", "FanInProd"]:
          if axis != 0:
              raise absltest.SkipTest(
                  "`FanInSum` and `FanInProd` are skipped when "
                  "axis != 0."
              )
          axis = None
      if (
          fan_in_mode == "FanInSum" or axis == 0
      ) and branch_in == "dense_after_branch_in":
          raise absltest.SkipTest(
              "`FanInSum` and `FanInConcat(0)` "
              "require `is_gaussian`."
          )
  
      if (
          axis == 1 or fan_in_mode == "FanInProd"
      ) and branch_in == "dense_before_branch_in":
          raise absltest.SkipTest(
              "`FanInConcat` or `FanInProd` on feature axis requires a dense layer "
              "after concatenation or Hadamard product."
          )
      if fan_in_mode == "FanInSum":
          fan_in_layer = stax.FanInSum()
      elif fan_in_mode == "FanInProd":
          fan_in_layer = stax.FanInProd()
      else:
          fan_in_layer = stax.FanInConcat(axis)
  
      if n_branches != 2:
          test_utils.skip_test(self)
  
      key = random.PRNGKey(1)
      X0_1 = jnp.cos(random.normal(key, (4, 3)))
      X0_2 = None if same_inputs else random.normal(key, (8, 3))
  
      width = 1024
      n_samples = 256 * 2
  
      if default_backend() == "tpu":
          tol = 0.07
      else:
          tol = 0.02
  
      dense = stax.Dense(width, 1.25, 0.1)
      input_layers = [dense, stax.FanOut(n_branches)]
  
      branches = []
      for b in range(n_branches):
          branch_layers = [FanInTest._get_phi(b)]
          for i in range(b):
              multiplier = 1 if axis not in (1, -1) else (1 + 0.25 * i)
              branch_layers += [
                  stax.Dense(
                      int(width * multiplier), 1.0 + 2 * i, 0.5 + i
                  ),
                  FanInTest._get_phi(i),
              ]
  
          if branch_in == "dense_before_branch_in":
              branch_layers += [dense]
          branches += [stax.serial(*branch_layers)]
  
      output_layers = [fan_in_layer, stax.Relu()]
      if branch_in == "dense_after_branch_in":
          output_layers.insert(1, dense)
  
      nn = stax.serial(
          *(input_layers + [stax.parallel(*branches)] + output_layers)
      )
  
      if get == "nngp":
          init_fn, apply_fn, kernel_fn = nn
      elif get == "ntk":
          init_fn, apply_fn, kernel_fn = stax.serial(
              nn, stax.Dense(1, 1.25, 0.5)
          )
      else:
          raise ValueError(get)
  
      kernel_fn_mc = nt.monte_carlo_kernel_fn(
          init_fn,
          apply_fn,
          key,
          n_samples,
          device_count=0 if axis in (0, -2) else -1,
          implementation=_DEFAULT_TESTING_NTK_IMPLEMENTATION,
          vmap_axes=None if axis in (0, -2) else 0,
      )
  
      exact = kernel_fn(X0_1, X0_2, get=get)
      empirical = kernel_fn_mc(X0_1, X0_2, get=get)
      test_utils.assert_close_matrices(self, empirical, exact, tol)
  
file_path: neural-tangents/tests/stax/branching_test.py
imports: import random
  from absl.testing import absltest
  from jax import config
  from jax import default_backend
  from jax import random
  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
label_desc: Use image node_id 2 for calling the FanInTest obj's underlying member method code with example usage: obj.test_fan_in_fc(same_inputs, axis, n_branches, get, branch_in, fan_in_mode) without return types