code: def _validate_elem_length(max_num_levels, elems_flat, axis):
      """Checks that elems all have the same length, and returns that length."""
      assertions = []
  
      elem_length = ps.shape(elems_flat[0])[axis]
  
      # The default size limit will overflow a 32-bit int, so make sure we're
      # using 64-bit.
      size_limit = 2 ** (ps.cast(max_num_levels, np.int64) + 1)
      enough_levels = ps.less(
          ps.cast(elem_length, np.int64), size_limit
      )
      enough_levels_ = tf.get_static_value(enough_levels)
      if enough_levels_ is None:
          assertions.append(
              tf.debugging.assert_equal(
                  enough_levels,
                  True,
                  message="Input `Tensor`s must have dimension less than"
                  " `2**(max_num_levels + 1)` along `axis=={}`."
                  " (saw: {} which is not less than 2**{} == {})".format(
                      axis, elem_length, max_num_levels, size_limit
                  ),
              )
          )
      elif not enough_levels_:
          raise ValueError(
              "Input `Tensor`s must have dimension less than"
              " `2**(max_num_levels + 1)` along `axis == {}`"
              " (saw: {} which is not less than 2**{} == {})".format(
                  axis, elem_length, max_num_levels, size_limit
              )
          )
  
      is_consistent = ps.reduce_all(
          [
              ps.equal(ps.shape(elem)[axis], elem_length)
              for elem in elems_flat[1:]
          ]
      )
  
      is_consistent_ = tf.get_static_value(is_consistent)
      if is_consistent_ is None:
          assertions.append(
              tf.debugging.assert_equal(
                  is_consistent,
                  True,
                  message="Inputs must have the same size along the given axis."
                  " (saw: {})".format(
                      [elem.shape for elem in elems_flat]
                  ),
              )
          )
      elif not is_consistent_:
          raise ValueError(
              "Inputs must have the same size along the given axis."
              " (saw: {})".format([elem.shape for elem in elems_flat])
          )
      return elem_length, assertions
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 3 for calling a global function with example usage: _validate_elem_length(max_num_levels, elems_flat, axis) and returns: elem_length,  assertions