code: def test_linear_head(self, inputs, expected_shape):
      rng = random.PRNGKey(0)
      output, _ = heads.LinearHead().init_with_output(
          rng, inputs, None, "", train=False
      )
      self.assertTupleEqual(output.shape, expected_shape)
  
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 1 for calling the HeadsTest obj's underlying member method code with example usage: obj.test_linear_head(inputs, expected_shape) without return types