code: def _get_phi(cls, i):
      return {0: stax.Relu(), 1: stax.Erf(), 2: stax.Abs()}[i % 3]
  
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 1 for calling the FanInTest obj's underlying member method code with example usage: obj._get_phi(cls, i) and returns: dict