code: def _assert_batch_shape_matches_weights(
      distribution, weights_shape, diststr
  ):
      """Checks that all parts of a distribution have the expected batch shape."""
      shapes = [weights_shape] + tf.nest.flatten(
          distribution.batch_shape_tensor()
      )
      static_shapes = [
          tf.get_static_value(ps.convert_to_shape_tensor(s))
          for s in shapes
      ]
      static_shapes_not_none = [
          s for s in static_shapes if s is not None
      ]
      static_shapes_match = all(
          [
              np.all(
                  a == b
              )  # Also need to check for rank mismatch (below).
              for (a, b) in zip(
                  static_shapes_not_none[1:],
                  static_shapes_not_none[:-1],
              )
          ]
      )
  
      # Build a separate list of static ranks, since rank is often static even when
      # shape is not.
      ranks = [ps.rank_from_shape(s) for s in shapes]
      static_ranks = [int(r) for r in ranks if not tf.is_tensor(r)]
      static_ranks_match = all(
          [
              a == b
              for (a, b) in zip(static_ranks[1:], static_ranks[:-1])
          ]
      )
  
      msg = (
          "The {diststr} distribution's batch shape does not match the particle "
          "weights; a correct {diststr} distribution must return an independent "
          "log-density for each particle. You may be "
          "creating a joint distribution in which some parts do not depend on the "
          "previous particles, and/or you are creating an autobatched joint "
          "distribution without setting `batch_ndims`.".format(
              diststr=diststr
          )
      )
      if not (static_ranks_match and static_shapes_match):
          raise ValueError(
              msg
              + " "
              + "Weights have shape {}, but the distribution has batch "
              "shape {}.".format(
                  weights_shape, distribution.batch_shape
              )
          )
  
      assertions = []
      if distribution.validate_args and any(
          [s is None for s in static_shapes]
      ):
          assertions = [
              assert_util.assert_equal(a, b, message=msg)
              for a, b in zip(shapes[1:], shapes[:-1])
          ]
      return assertions
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/particle_filter.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import sequential_monte_carlo_kernel
  from tensorflow_probability.python.experimental.mcmc import weighted_resampling
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import docstring_util
  from tensorflow_probability.python.internal import loop_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 9 for calling a global function with example usage: _assert_batch_shape_matches_weights(distribution, weights_shape, diststr) and returns: assertions