code: def _filter_stuck_chains(params):
      """Rough heuristic for stuck MCMC parameters.
  
      1. Compute the z scores of the noise_scale `variances`.
      2. Compute the z score of 0.
      3. Filter parameters with z score more than halfway to the score for 0.
  
      If there are 0 or 1 chains left, just return the two with the biggest
      variances. These might be stuck!
  
      Args:
        params: Nested dictionary with a `noise_scale` key.
  
      Returns:
        A dictionary with the same structure, but only leading dimensions with
        a reasonable amount of variance in the noise parameter.
      """
      # TODO(colcarroll): Use a better heuristic here for filtering stuck chains.
      if "noise_scale" not in params["params"]:
          return params
      stds = jnp.std(params["params"]["noise_scale"].squeeze(), axis=1)
      stds_mu, stds_scale = jnp.mean(stds), jnp.std(stds)
  
      z_scores = (stds - stds_mu) / stds_scale
      halfway_to_zero = -0.5 * stds_mu / stds_scale
      unstuck = jnp.where(z_scores > halfway_to_zero)[0]
      if unstuck.shape[0] > 2:
          return jax.tree_map(lambda x: x[unstuck], params)
      best_two = jnp.argsort(stds)[-2:]
      return jax.tree_map(lambda x: x[best_two], params)
  
imports: import functools
  from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
  import bayeux
  import jax
  import jax.numpy
  from jaxtyping import PyTree
  import matplotlib.pyplot
  import numpy
  import pandas
  from tensorflow_probability.python.experimental.autobnn import bnn
  from tensorflow_probability.python.experimental.autobnn import util
  from tensorflow_probability.python.experimental.timeseries import metrics
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 3 for calling a global function with example usage: _filter_stuck_chains(params) and returns: jax, params, jax