code: def test_fan_in_conv(
      self,
      same_inputs,
      axis,
      n_branches,
      get,
      branch_in,
      readout,
      fan_in_mode,
  ):
      test_utils.skip_test(self)
      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 in [0, 1, 2]
      ) and branch_in == "dense_after_branch_in":
          raise absltest.SkipTest(
              "`FanInSum` and `FanInConcat(0/1/2)` "
              "require `is_gaussian`."
          )
  
      if (
          axis == 3 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)
  
      key = random.PRNGKey(1)
      X0_1 = random.normal(key, (2, 5, 6, 3))
      X0_2 = None if same_inputs else random.normal(key, (3, 5, 6, 3))
  
      if default_backend() == "tpu":
          width = 2048
          n_samples = 1024
          tol = 0.02
      else:
          width = 1024
          n_samples = 512
          tol = 0.015
  
      conv = stax.Conv(
          out_chan=width,
          filter_shape=(3, 3),
          padding="SAME",
          W_std=1.25,
          b_std=0.1,
      )
  
      input_layers = [conv, 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 (3, -1) else (1 + 0.25 * i)
              branch_layers += [
                  stax.Conv(
                      out_chan=int(width * multiplier),
                      filter_shape=(i + 1, 4 - i),
                      padding="SAME",
                      W_std=1.25 + i,
                      b_std=0.1 + i,
                  ),
                  FanInTest._get_phi(i),
              ]
  
          if branch_in == "dense_before_branch_in":
              branch_layers += [conv]
          branches += [stax.serial(*branch_layers)]
  
      output_layers = [
          fan_in_layer,
          stax.Relu(),
          stax.GlobalAvgPool() if readout == "pool" else stax.Flatten(),
      ]
      if branch_in == "dense_after_branch_in":
          output_layers.insert(1, conv)
  
      nn = stax.serial(
          *(input_layers + [stax.parallel(*branches)] + output_layers)
      )
  
      init_fn, apply_fn, kernel_fn = stax.serial(
          nn, stax.Dense(1 if get == "ntk" else width, 1.25, 0.5)
      )
  
      kernel_fn_mc = nt.monte_carlo_kernel_fn(
          init_fn,
          apply_fn,
          key,
          n_samples,
          device_count=0 if axis in (0, -4) else -1,
          implementation=_DEFAULT_TESTING_NTK_IMPLEMENTATION,
          vmap_axes=None if axis in (0, -4) 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 3 for calling the FanInTest obj's underlying member method code with example usage: obj.test_fan_in_conv(same_inputs, axis, n_branches, get, branch_in, readout, fan_in_mode) without return types