code: def main(unused_argv):
      x1 = tf.random.normal((6, 8, 8, 3), seed=1)
      x2 = tf.random.normal((3, 8, 8, 3), seed=2)
  
      # A vanilla CNN `tf.keras.Model` example.
      print("A Keras CNN example.")
  
      f = tf.keras.Sequential()
      f.add(tf.keras.layers.Conv2D(16, (3, 3), activation="relu"))
      f.add(tf.keras.layers.Conv2D(16, (3, 3), activation="relu"))
      f.add(tf.keras.layers.Conv2D(16, (3, 3)))
      f.add(tf.keras.layers.Flatten())
      f.add(tf.keras.layers.Dense(10))
  
      f.build((None, *x1.shape[1:]))
  
      _, params = nt.experimental.get_apply_fn_and_params(f)
      _compute_and_check_ntks(f, x1, x2, params)
  
      # A `tf.function` example.
      print("A `tf.function` example.")
  
      params_tf = tf.random.normal((1, 2, 3, 4), seed=3)
  
      @tf.function(
          input_signature=[
              tf.TensorSpec(None),
              tf.TensorSpec((None, *x1.shape[1:])),
          ]
      )
      def f_tf(params, x):
          return (
              tf.transpose(x, (0, 3, 1, 2))
              * tf.reduce_mean(params**2)
              + 1.0
          )
  
      _compute_and_check_ntks(f_tf, x1, x2, params_tf)
  
imports: from absl import app
  import neural_tangents
  import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: main(unused_argv) and returns: unknown