code: def testRandnInitShape(self, shape):
      key = random.PRNGKey(0)
      out = stax.randn()(key, shape)
      self.assertEqual(out.shape, shape)
  
imports: from absl.testing import absltest
  import numpy
  from jax._src import test_util
  from jax import random
  from jax.example_libraries import stax
  from jax import dtypes
  from jax import config
label_desc: Use image node_id 1 for calling the StaxTest obj's underlying member method code with example usage: obj.testRandnInitShape(shape) without return types