code: def test_localization_head_tal(
      self,
      x,
      feature_pyramid_config,
      weight_sharing,
      output_per_class_displacements,
      expected_output_shape,
      expected_keys,
  ):
      rng = random.PRNGKey(0)
      output, params = heads.LocalizationHead(
          num_conv_layers=3,
          kernel_size=3,
          num_classes=10,
          feature_pyramid_config=feature_pyramid_config,
          weight_sharing=weight_sharing,
          output_per_class_displacements=output_per_class_displacements,
      ).init_with_output(
          rng,
          x,
          None,
          "temporal_localization",
          train=False,
      )
      self.assertTupleEqual(output.shape, expected_output_shape)
      self.assertSetEqual(set(params["params"].keys()), expected_keys)
  
file_path: scenic/scenic/projects/unloc/heads_test.py
imports: from absl.testing import parameterized
  from jax import random
  import ml_collections
  import numpy
  from scenic.projects.unloc import heads
  import tensorflow
label_desc: Use image node_id 3 for calling the HeadsTest obj's underlying member method code with example usage: obj.test_localization_head_tal(x, feature_pyramid_config, weight_sharing, output_per_class_displacements, expected_output_shape, expected_keys) without return types