code: def test__filter_stuck_chains_doesnt_overfilter(self):
      noise_scale = 0.001 * np.random.randn(64, 100, 1)
      filtered = training_util._filter_stuck_chains(
          {"params": {"noise_scale": noise_scale}}
      )
      self.assertEqual(
          filtered["params"]["noise_scale"].shape, (64, 100, 1)
      )
  
imports: import chex
  import jax
  import jax.numpy
  import numpy
  from tensorflow_probability.python.experimental.autobnn import kernels
  from tensorflow_probability.python.experimental.autobnn import operators
  from tensorflow_probability.python.experimental.autobnn import training_util
  from tensorflow_probability.python.experimental.autobnn import util
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the TrainingUtilTest obj's underlying member method code with example usage: obj.test__filter_stuck_chains_doesnt_overfilter() without return types