code: def test_keras_functional(
      self,
      f,
      input_shape,
      trace_axes,
      diagonal_axes,
      vmap_axes,
  ):
      f = f(classes=1, input_shape=input_shape, weights=None)
      f.build((None, *input_shape))
      f_jax, params = experimental.get_apply_fn_and_params(f)
      self._compare_ntks(
          f, f_jax, params, trace_axes, diagonal_axes, vmap_axes
      )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import numpy
  import neural_tangents
  from neural_tangents import experimental
  import numpy
  import tensorflow
label_desc: Use image node_id 2 for calling the EmpiricalTfTest obj's underlying member method code with example usage: obj.test_keras_functional(f, input_shape, trace_axes, diagonal_axes, vmap_axes) without return types