docstr: Tests https://github.com/google/neural-tangents/issues/123.
code: def test_issue_123(self, get, architecture, do_jit):
      """Tests https://github.com/google/neural-tangents/issues/123."""
      if architecture == "wrn":
          # https://github.com/google/neural-tangents/issues/123#issue-992927376
          def WideResnetBlock(
              channels, strides=(1, 1), channel_mismatch=False
          ):
              main = stax.serial(
                  stax.Relu(),
                  stax.Conv(
                      channels,
                      (3, 3),
                      strides,
                      padding="SAME",
                      parameterization="standard",
                  ),
                  stax.Relu(),
                  stax.Conv(
                      channels,
                      (3, 3),
                      padding="SAME",
                      parameterization="standard",
                  ),
              )
              shortcut = (
                  stax.Identity()
                  if not channel_mismatch
                  else stax.Conv(
                      channels,
                      (3, 3),
                      strides,
                      padding="SAME",
                      parameterization="standard",
                  )
              )
              return stax.serial(
                  stax.FanOut(2),
                  stax.parallel(main, shortcut),
                  stax.FanInSum(),
              )
  
          def WideResnetGroup(n, channels, strides=(1, 1)):
              blocks = []
              blocks += [
                  WideResnetBlock(
                      channels, strides, channel_mismatch=True
                  )
              ]
              for _ in range(n - 1):
                  blocks += [WideResnetBlock(channels, (1, 1))]
              return stax.serial(*blocks)
  
          def WideResnet(block_size, k, num_classes):
              return stax.serial(
                  stax.Conv(
                      16,
                      (3, 3),
                      padding="SAME",
                      parameterization="standard",
                  ),
                  WideResnetGroup(block_size, int(16 * k)),
                  WideResnetGroup(block_size, int(32 * k), (2, 2)),
                  WideResnetGroup(block_size, int(64 * k), (2, 2)),
                  stax.AvgPool((8, 8), padding="SAME"),
                  stax.Flatten(),
                  stax.Dense(
                      num_classes, 1.0, 0.0, parameterization="standard"
                  ),
              )
  
          init_fn, apply_fn, kernel_fn = WideResnet(
              block_size=1, k=1, num_classes=1
          )
  
      elif architecture == "conv":
          # https://github.com/google/neural-tangents/issues/123#issuecomment-932809224
          init_fn, apply_fn, kernel_fn = stax.serial(
              stax.Conv(1, (3, 3)),
              stax.Relu(),
              stax.Flatten(),
          )
  
      else:
          raise ValueError(architecture)
  
      x1 = x2 = jnp.zeros((1, 8, 8, 3))
  
      def kernel_scalar(x1, x2):
          return kernel_fn(x1, x2, get)[0, 0]
  
      if do_jit:
          kernel_scalar = jit(kernel_scalar)
  
      # Compare forward pass to `value_and_grad`.
      k1 = kernel_scalar(x1, x2)
      k2, k2_grad = value_and_grad(kernel_scalar)(x1, x2)
      self.assertAllClose(k1, k2)
  
      # Compare to forward-mode.
      k2_fwd, _ = jvp(kernel_scalar, (x1, x2), (x1, x2))
      k2_grad_fwd = jacfwd(kernel_scalar)(x1, x2)
      self.assertAllClose(k1, k2_fwd)
      self.assertAllClose(k2_grad, k2_grad_fwd)
  
      # Compare to 0.
      self.assertAllClose(
          grad(kernel_scalar)(x1, x2), jnp.zeros_like(x1)
      )
  
imports: import itertools
  import random
  from absl.testing import absltest
  from jax import config
  from jax import default_backend
  from jax import grad, jacfwd, jacrev, jit, jvp, value_and_grad
  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 AutodiffTest obj's underlying member method code with example usage: obj.test_issue_123(get, architecture, do_jit) and returns: kernel_fn, stax, stax, stax