code: def test_dropout_broadcasts(self):
      x_in = jnp.ones([3, 3, 3, 3])
      x_out = basic.dropout(
          base.next_rng_key(), 0.5, x_in, broadcast_dims=(1,)
      )
      np.testing.assert_allclose(x_out[:, 0], x_out[:, 1])
      np.testing.assert_allclose(x_out[:, 0], x_out[:, 2])
      self.assertRaises(
          AssertionError,
          np.testing.assert_allclose,
          x_out[:, :, 0],
          x_out[:, :, 1],
      )
      self.assertRaises(
          AssertionError,
          np.testing.assert_allclose,
          x_out[:, :, 0],
          x_out[:, :, 2],
      )
      self.assertRaises(
          AssertionError, np.testing.assert_allclose, x_out[0], x_out[1]
      )
      self.assertEqual(x_in.shape, x_out.shape)
  
      x_out = basic.dropout(
          base.next_rng_key(), 0.5, x_in, broadcast_dims=(1, 2)
      )
      np.testing.assert_allclose(x_out[:, 0], x_out[:, 1])
      np.testing.assert_allclose(x_out[:, 0], x_out[:, 2])
      np.testing.assert_allclose(x_out[:, :, 0], x_out[:, :, 1])
      np.testing.assert_allclose(x_out[:, :, 0], x_out[:, :, 2])
      self.assertRaises(
          AssertionError, np.testing.assert_allclose, x_out[0], x_out[1]
      )
      self.assertEqual(x_in.shape, x_out.shape)
  
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 7 for calling the BasicTest obj's underlying member method code with example usage: obj.test_dropout_broadcasts() without return types