tensorflow_probability/python/internal/distribution_util.py
def shapes_from_loc_and_scale(loc, scale, name="shapes_from_loc_and_scale"):
  """Infer distribution batch and event shapes from a location and scale.

  Location and scale family distributions determine their batch/event shape by
  broadcasting the `loc` and `scale` args.  This helper does that broadcast,
  statically if possible.

  Batch shape broadcasts as per the normal rules.
  We allow the `loc` event shape to broadcast up to that of `scale`.  We do not
  allow `scale`'s event shape to change.  Therefore, the last dimension of `loc`
  must either be size `1`, or the same as `scale.range_dimension`.

  See `MultivariateNormalLinearOperator` for a usage example.

  Args:
    loc: `Tensor` (already converted to tensor) or `None`. If `None`, or
      `rank(loc)==0`, both batch and event shape are determined by `scale`.
    scale:  A `LinearOperator` instance.
    name:  A string name to prepend to created ops.

  Returns:
    batch_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.
    event_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.

  Raises:
    ValueError:  If the last dimension of `loc` is determined statically to be
      different than the range of `scale`.
  """
  if loc is not None and tensorshape_util.rank(loc.shape) == 0:
    loc = None  # scalar loc is irrelevant to determining batch/event shape.
  with tf.name_scope(name):
    # Get event shape.
    event_size = scale.range_dimension_tensor()
    event_size_ = tf.get_static_value(event_size)
    loc_event_size_ = (None if loc is None
                       else tf.compat.dimension_value(loc.shape[-1]))

    if event_size_ is not None and loc_event_size_ is not None:
      # Static check that event shapes match.
      if loc_event_size_ != 1 and loc_event_size_ != event_size_:
        raise ValueError(
            "Event size of 'scale' ({}) could not be broadcast up to that "
            "of 'loc' ({}).".format(event_size_, loc_event_size_))
    elif loc_event_size_ is not None and loc_event_size_ != 1:
      event_size_ = loc_event_size_

    if event_size_ is None:
      event_shape = event_size[tf.newaxis]
    else:
      event_shape = tf.convert_to_tensor(
          value=np.reshape(event_size_, [1]),
          dtype=tf.int32,
          name="event_shape")

    # Get batch shape.
    batch_shape = scale.batch_shape_tensor()
    if loc is not None:
      loc_batch_shape = tensorshape_util.with_rank_at_least(loc.shape, 1)[:-1]
      if tensorshape_util.rank(
          loc.shape) is None or not tensorshape_util.is_fully_defined(
              loc_batch_shape):
        loc_batch_shape = tf.shape(input=loc)[:-1]
      else:
        loc_batch_shape = tf.convert_to_tensor(
            value=loc_batch_shape, dtype=tf.int32, name="loc_batch_shape")
      # This is defined in the core util module.
      batch_shape = prefer_static_broadcast_shape(batch_shape, loc_batch_shape)  # pylint: disable=undefined-variable
      batch_shape = tf.convert_to_tensor(
          value=batch_shape, dtype=tf.int32, name="batch_shape")

    return batch_shape, event_shape
def shapes_from_loc_and_scale(loc, scale, name="shapes_from_loc_and_scale"):
  """Infer distribution batch and event shapes from a location and scale.

  Location and scale family distributions determine their batch/event shape by
  broadcasting the `loc` and `scale` args.  This helper does that broadcast,
  statically if possible.

  Batch shape broadcasts as per the normal rules.
  We allow the `loc` event shape to broadcast up to that of `scale`.  We do not
  allow `scale`'s event shape to change.  Therefore, the last dimension of `loc`
  must either be size `1`, or the same as `scale.range_dimension`.

  See `MultivariateNormalLinearOperator` for a usage example.

  Args:
    loc: `Tensor` (already converted to tensor) or `None`. If `None`, or
      `rank(loc)==0`, both batch and event shape are determined by `scale`.
    scale:  A `LinearOperator` instance.
    name:  A string name to prepend to created ops.

  Returns:
    batch_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.
    event_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.

  Raises:
    ValueError:  If the last dimension of `loc` is determined statically to be
      different than the range of `scale`.
  """
  if loc is not None and tensorshape_util.rank(loc.shape) == 0:
    loc = None  # scalar loc is irrelevant to determining batch/event shape.
  with tf.name_scope(name):
    # Get event shape.
    event_size = scale.range_dimension_tensor()
    event_size_ = tf.get_static_value(event_size)
    loc_event_size_ = (None if loc is None
                       else tf.compat.dimension_value(loc.shape[-1]))

    if event_size_ is not None and loc_event_size_ is not None:
      # Static check that event shapes match.
      if loc_event_size_ != 1 and loc_event_size_ != event_size_:
        raise ValueError(
            "Event size of 'scale' ({}) could not be broadcast up to that "
            "of 'loc' ({}).".format(event_size_, loc_event_size_))
    elif loc_event_size_ is not None and loc_event_size_ != 1:
      event_size_ = loc_event_size_

    if event_size_ is None:
      event_shape = event_size[tf.newaxis]
    else:
      event_shape = tf.convert_to_tensor(
          value=np.reshape(event_size_, [1]),
          dtype=tf.int32,
          name="event_shape")

    # Get batch shape.
    batch_shape = scale.batch_shape_tensor()
    if loc is not None:
      loc_batch_shape = tensorshape_util.with_rank_at_least(loc.shape, 1)[:-1]
      if tensorshape_util.rank(
          loc.shape) is None or not tensorshape_util.is_fully_defined(
              loc_batch_shape):
        loc_batch_shape = tf.shape(input=loc)[:-1]
      else:
        loc_batch_shape = tf.convert_to_tensor(
            value=loc_batch_shape, dtype=tf.int32, name="loc_batch_shape")
      # This is defined in the core util module.
      batch_shape = prefer_static_broadcast_shape(batch_shape, loc_batch_shape)  # pylint: disable=undefined-variable
      batch_shape = tf.convert_to_tensor(
          value=batch_shape, dtype=tf.int32, name="batch_shape")

    return batch_shape, event_shape
Infer distribution batch and event shapes from a location and scale.

  Location and scale family distributions determine their batch/event shape by
  broadcasting the `loc` and `scale` args.  This helper does that broadcast,
  statically if possible.

  Batch shape broadcasts as per the normal rules.
  We allow the `loc` event shape to broadcast up to that of `scale`.  We do not
  allow `scale`'s event shape to change.  Therefore, the last dimension of `loc`
  must either be size `1`, or the same as `scale.range_dimension`.

  See `MultivariateNormalLinearOperator` for a usage example.

  Args:
    loc: `Tensor` (already converted to tensor) or `None`. If `None`, or
      `rank(loc)==0`, both batch and event shape are determined by `scale`.
    scale:  A `LinearOperator` instance.
    name:  A string name to prepend to created ops.

  Returns:
    batch_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.
    event_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.

  Raises:
    ValueError:  If the last dimension of `loc` is determined statically to be
      different than the range of `scale`.