code: def test_multinomial_r1(self, num_samples):
      out = basic.multinomial(
          random.PRNGKey(428), jnp.ones([4]), num_samples
      )
      self.assertEqual(out.shape, (num_samples,))
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import base
  from haiku._src import basic
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  from jax import random
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the BasicTest obj's underlying member method code with example usage: obj.test_multinomial_r1(num_samples) without return types