code: def test_linear():
      from deepchem.models.jax_models import layers as jax_models_layers
  
      def forward(x):
          layer = jax_models_layers.Linear(2)
          return layer(x)
  
      forward = hk.transform(forward)
      rng = jax.random.PRNGKey(42)
      x = jnp.ones([8, 28 * 28])
      params = forward.init(rng, x)
      output = forward.apply(params, rng, x)
      assert output.shape == (8, 2)
  
label_desc: Use image node_id 1 for calling a global function with example usage: test_linear() and returns: layer