code: def _interleave(a, b, axis):
      """Interleaves two `Tensor`s along the given axis."""
      # [a b c ...] [d e f ...] -> [a d b e c f ...]
      num_elems_a = ps.shape(a)[axis]
      num_elems_b = ps.shape(b)[axis]
  
      # Note that interleaving implies rank(a)==rank(b).
      axis = ps.where(axis >= 0, axis, ps.rank(a) + axis)
      axis = (
          int(axis)  # Avoid ndarray values.
          if tf.get_static_value(axis) is not None
          else axis
      )
  
      def _interleave_with_b(a):
          return tf.reshape(
              # Work around lack of support for Tensor axes in `tf.stack` by using
              # `concat` and `expand_dims` instead.
              tf.concat(
                  [
                      tf.expand_dims(a, axis=axis + 1),
                      tf.expand_dims(b, axis=axis + 1),
                  ],
                  axis=axis + 1,
              ),
              ps.concat(
                  [
                      ps.shape(a)[:axis],
                      [2 * num_elems_b],
                      ps.shape(a)[axis + 1 :],
                  ],
                  axis=0,
              ),
          )
  
      return ps.cond(
          ps.equal(num_elems_a, num_elems_b + 1),
          lambda: tf.concat(
              [  # pylint: disable=g-long-lambda
                  _interleave_with_b(
                      _slice_along_axis(a, None, -1, axis=axis)
                  ),
                  _slice_along_axis(a, -1, None, axis=axis),
              ],
              axis=axis,
          ),
          lambda: _interleave_with_b(a),
      )
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 2 for calling a global function with example usage: _interleave(a, b, axis) and returns: ps, tf