MixturesSameFamily layer when data format is channels_first

Can someone please tell me how to change the source code in class MixtureSameFamily(DistributionLambda) in distribution layer when the data format is channels_first? https://github.com/tensorflow/probability/blob/9e50fa39f37216dd084c24aa3dd14ab0d4dee926/tensorflow_probability/python/layers/distribution_layer.py#L1396-L1520

I've been trying to figure it out for a week but no luck.

The errors seem to be occurring from this function:

 @staticmethod
  def new(params, num_components, component_layer,
          validate_args=False, name=None):
    """Create the distribution instance from a `params` vector."""
    with tf.name_scope(name or 'MixtureSameFamily'):
      params = tf.convert_to_tensor(params, name='params')
      num_components = tf.convert_to_tensor(
          num_components, name='num_components', dtype_hint=tf.int32)

      components_dist = component_layer(
          tf.reshape(
              params[..., num_components:],
              tf.concat([tf.shape(params)[:-1], [num_components, -1]],
                        axis=0)))
      mixture_dist = categorical_lib.Categorical(
          logits=params[..., :num_components])
      return mixture_same_family_lib.MixtureSameFamily(
          mixture_dist,
          components_dist,
          # TODO(b/120154797): Change following to `validate_args=True` after
          # fixing: "ValueError: `mixture_distribution` must have scalar
          # `event_dim`s." assertion in MixtureSameFamily.
          validate_args=False)

Start by reading the MixtureSameFamily class and its new method in tensorflow_probability/python/layers/distribution_layer.py at the linked lines. Reproduce the reported errors with channels_first and compare the parameter reshaping and component-layer inputs with the expected data layout. Done means the layer handles channels_first without the reported errors and the relevant behavior is verified.