tensorflow_probability/python/layers/distribution_layer.py
def new(params, num_components, component_layer,
          validate_args=False, name=None):
    """Create the distribution instance from a `params` vector."""
    with tf.compat.v1.name_scope(name, 'MixtureSameFamily',
                                 [params, num_components, component_layer]):
      params = tf.convert_to_tensor(value=params, name='params')
      num_components = tf.convert_to_tensor(
          value=num_components, name='num_components', dtype_hint=tf.int32)

      components_dist = component_layer(
          tf.reshape(
              params[..., num_components:],
              tf.concat([tf.shape(input=params)[:-1], [num_components, -1]],
                        axis=0)))
      mixture_dist = tfd.Categorical(logits=params[..., :num_components])
      return tfd.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)
def new(params, num_components, component_layer,
          validate_args=False, name=None):
    """Create the distribution instance from a `params` vector."""
    with tf.compat.v1.name_scope(name, 'MixtureSameFamily',
                                 [params, num_components, component_layer]):
      params = tf.convert_to_tensor(value=params, name='params')
      num_components = tf.convert_to_tensor(
          value=num_components, name='num_components', dtype_hint=tf.int32)

      components_dist = component_layer(
          tf.reshape(
              params[..., num_components:],
              tf.concat([tf.shape(input=params)[:-1], [num_components, -1]],
                        axis=0)))
      mixture_dist = tfd.Categorical(logits=params[..., :num_components])
      return tfd.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)
Create the distribution instance from a `params` vector.