code: def get_apply_fn_and_params(f: tf.Module):
      """Converts a :class:`tf.Module` into a forward-pass `apply_fn` and `params`.
  
      Use this function to extract `params` to pass to the Tensorflow empirical NTK
      kernel function.
  
      .. warning::
        This function does not distinguish between trainable and non-trainable
        parameters of the model.
  
      Args:
        f:
          a :class:`tf.Module` to convert to a `apply_fn(params, x)` function. Must
          have an `input_shape` attribute set (specifying shape of `x`), and be
          callable or be a :class:`tf.keras.Model`.
  
      Returns:
        A tuple fo `(apply_fn, params)`, where `params` is a `PyTree[tf.Tensor]`.
      """
  
      @tf.function
      def forward_tf(x: PyTree) -> PyTree:
          if isinstance(f, tf.keras.Model):
              return f.call(x, training=False)
  
          if not hasattr(f, "__call__"):
              raise NotImplementedError(
                  f"Got `f={f}` of type {type(f)}, "
                  f"that is not callable. Please file a bug at "
                  f"https://github.com/google/neural-tangents."
              )
  
          return f(x)
  
      if not hasattr(f, "input_shape"):
          raise NotImplementedError(
              f"`f={f}` must have `input_shape` set. "
              f"Please file a bug at "
              f"https://github.com/google/neural-tangents."
          )
  
      apply_fn_, params = tf2jax.convert(
          forward_tf, tf.TensorSpec(f.input_shape)
      )
  
      def apply_fn(params: PyTree, x: PyTree) -> PyTree:
          outputs, _ = apply_fn_(
              params, x
          )  # Dropping parameters (not updated).
          return outputs
  
      return apply_fn, params
  
file_path: neural-tangents/neural_tangents/experimental/empirical_tf/empirical.py
imports: from typing import Callable, Optional, Union
  import warnings
  from jax.experimental import jax2tf
  from neural_tangents._src.empirical import _DEFAULT_NTK_FWD
  from neural_tangents._src.empirical import _DEFAULT_NTK_J_RULES
  from neural_tangents._src.empirical import _DEFAULT_NTK_S_RULES
  from neural_tangents._src.empirical import DEFAULT_NTK_IMPLEMENTATION
  from neural_tangents._src.empirical import empirical_ntk_fn
  from neural_tangents._src.empirical import NtkImplementation
  from neural_tangents._src.utils.typing import Axes
  from neural_tangents._src.utils.typing import PyTree
  from neural_tangents._src.utils.typing import VMapAxes
  import tensorflow
  import tf2jax
label_desc: Use image node_id 2 for calling a global function with example usage: get_apply_fn_and_params(f) and returns: apply_fn,  params, f, outputs, f