tensorflow_probability/python/distributions/linear_gaussian_ssm.py
def _augment_sample_shape(partial_batch_dist,
                          full_sample_and_batch_shape,
                          validate_args=False):
  """Augment a sample shape to broadcast batch dimensions.

  Computes an augmented sample shape, so that any batch dimensions not
  part of the distribution `partial_batch_dist` are treated as identical
  distributions.

  # partial_batch_dist.batch_shape  = [      7]
  # full_sample_and_batch_shape     = [3, 4, 7]
  # => return an augmented sample shape of [3, 4] so that
  #    partial_batch_dist.sample(augmented_sample_shape) has combined
  #    sample and batch shape of [3, 4, 7].

  Args:
    partial_batch_dist: `tfd.Distribution` instance with batch shape a
      prefix of `full_sample_and_batch_shape`.
    full_sample_and_batch_shape: a Tensor or Tensor-like shape.
    validate_args: if True, check for shape errors at runtime.
  Returns:
    augmented_sample_shape: sample shape such that
      `partial_batch_dist.sample(augmented_sample_shape)` has combined
      sample and batch shape of `full_sample_and_batch_shape`.

  Raises:
    ValueError: if `partial_batch_dist.batch_shape` has more dimensions than
      `full_sample_and_batch_shape`.
    NotImplementedError: if broadcasting would be required to make
      `partial_batch_dist.batch_shape` into a prefix of
      `full_sample_and_batch_shape` .
  """
  full_ndims = distribution_util.prefer_static_shape(
      full_sample_and_batch_shape)[0]
  partial_batch_ndims = (
      tensorshape_util.rank(partial_batch_dist.batch_shape)  # pylint: disable=g-long-ternary
      if tensorshape_util.rank(partial_batch_dist.batch_shape) is not None
      else distribution_util.prefer_static_shape(
          partial_batch_dist.batch_shape_tensor())[0])

  num_broadcast_dims = full_ndims - partial_batch_ndims

  expected_partial_batch_shape = (
      full_sample_and_batch_shape[num_broadcast_dims:])
  expected_partial_batch_shape_static = tf.get_static_value(
      full_sample_and_batch_shape[num_broadcast_dims:])

  # Raise errors statically if possible.
  num_broadcast_dims_static = tf.get_static_value(num_broadcast_dims)
  if num_broadcast_dims_static is not None:
    if num_broadcast_dims_static < 0:
      raise ValueError("Cannot broadcast distribution {} batch shape to "
                       "target batch shape with fewer dimensions"
                       .format(partial_batch_dist))
  if (expected_partial_batch_shape_static is not None and
      tensorshape_util.is_fully_defined(partial_batch_dist.batch_shape)):
    if (partial_batch_dist.batch_shape and
        any(expected_partial_batch_shape_static != tensorshape_util.as_list(
            partial_batch_dist.batch_shape))):
      raise NotImplementedError("Broadcasting is not supported; "
                                "unexpected batch shape "
                                "(expected {}, saw {}).".format(
                                    expected_partial_batch_shape_static,
                                    partial_batch_dist.batch_shape
                                ))
  runtime_assertions = []
  if validate_args:
    runtime_assertions.append(
        assert_util.assert_greater_equal(
            tf.convert_to_tensor(value=num_broadcast_dims, dtype=tf.int32),
            tf.zeros((), dtype=tf.int32),
            message=("Cannot broadcast distribution {} batch shape to "
                     "target batch shape with fewer dimensions.".format(
                         partial_batch_dist))))
    runtime_assertions.append(
        assert_util.assert_equal(
            expected_partial_batch_shape,
            partial_batch_dist.batch_shape_tensor(),
            message=("Broadcasting is not supported; "
                     "unexpected batch shape."),
            name="assert_batch_shape_same"))

  with tf.control_dependencies(runtime_assertions):
    return full_sample_and_batch_shape[:num_broadcast_dims]
def _augment_sample_shape(partial_batch_dist,
                          full_sample_and_batch_shape,
                          validate_args=False):
  full_ndims = distribution_util.prefer_static_shape(
      full_sample_and_batch_shape)[0]
  partial_batch_ndims = (
      tensorshape_util.rank(partial_batch_dist.batch_shape)  # pylint: disable=g-long-ternary
      if tensorshape_util.rank(partial_batch_dist.batch_shape) is not None
      else distribution_util.prefer_static_shape(
          partial_batch_dist.batch_shape_tensor())[0])

  num_broadcast_dims = full_ndims - partial_batch_ndims

  expected_partial_batch_shape = (
      full_sample_and_batch_shape[num_broadcast_dims:])
  expected_partial_batch_shape_static = tf.get_static_value(
      full_sample_and_batch_shape[num_broadcast_dims:])

  # Raise errors statically if possible.
  num_broadcast_dims_static = tf.get_static_value(num_broadcast_dims)
  if num_broadcast_dims_static is not None:
    if num_broadcast_dims_static < 0:
      raise ValueError("Cannot broadcast distribution {} batch shape to "
                       "target batch shape with fewer dimensions"
                       .format(partial_batch_dist))
  if (expected_partial_batch_shape_static is not None and
      tensorshape_util.is_fully_defined(partial_batch_dist.batch_shape)):
    if (partial_batch_dist.batch_shape and
        any(expected_partial_batch_shape_static != tensorshape_util.as_list(
            partial_batch_dist.batch_shape))):
      raise NotImplementedError("Broadcasting is not supported; "
                                "unexpected batch shape "
                                "(expected {}, saw {}).".format(
                                    expected_partial_batch_shape_static,
                                    partial_batch_dist.batch_shape
                                ))
  runtime_assertions = []
  if validate_args:
    runtime_assertions.append(
        assert_util.assert_greater_equal(
            tf.convert_to_tensor(value=num_broadcast_dims, dtype=tf.int32),
            tf.zeros((), dtype=tf.int32),
            message=("Cannot broadcast distribution {} batch shape to "
                     "target batch shape with fewer dimensions.".format(
                         partial_batch_dist))))
    runtime_assertions.append(
        assert_util.assert_equal(
            expected_partial_batch_shape,
            partial_batch_dist.batch_shape_tensor(),
            message=("Broadcasting is not supported; "
                     "unexpected batch shape."),
            name="assert_batch_shape_same"))

  with tf.control_dependencies(runtime_assertions):
    return full_sample_and_batch_shape[:num_broadcast_dims]
def _augment_sample_shape ( partial_batch_dist , full_sample_and_batch_shape , validate_args = False ) : full_ndims = distribution_util . prefer_static_shape ( full_sample_and_batch_shape ) [ 0 ] partial_batch_ndims = ( tensorshape_util . rank ( partial_batch_dist . batch_shape ) # pylint: disable=g-long-ternary if tensorshape_util . rank ( partial_batch_dist . batch_shape ) is not None else distribution_util . prefer_static_shape ( partial_batch_dist . batch_shape_tensor ( ) ) [ 0 ] ) num_broadcast_dims = full_ndims - partial_batch_ndims expected_partial_batch_shape = ( full_sample_and_batch_shape [ num_broadcast_dims : ] ) expected_partial_batch_shape_static = tf . get_static_value ( full_sample_and_batch_shape [ num_broadcast_dims : ] ) # Raise errors statically if possible. num_broadcast_dims_static = tf . get_static_value ( num_broadcast_dims ) if num_broadcast_dims_static is not None : if num_broadcast_dims_static < 0 : raise ValueError ( "Cannot broadcast distribution {} batch shape to " "target batch shape with fewer dimensions" . format ( partial_batch_dist ) ) if ( expected_partial_batch_shape_static is not None and tensorshape_util . is_fully_defined ( partial_batch_dist . batch_shape ) ) : if ( partial_batch_dist . batch_shape and any ( expected_partial_batch_shape_static != tensorshape_util . as_list ( partial_batch_dist . batch_shape ) ) ) : raise NotImplementedError ( "Broadcasting is not supported; " "unexpected batch shape " "(expected {}, saw {})." . format ( expected_partial_batch_shape_static , partial_batch_dist . batch_shape ) ) runtime_assertions = [ ] if validate_args : runtime_assertions . append ( assert_util . assert_greater_equal ( tf . convert_to_tensor ( value = num_broadcast_dims , dtype = tf . int32 ) , tf . zeros ( ( ) , dtype = tf . int32 ) , message = ( "Cannot broadcast distribution {} batch shape to " "target batch shape with fewer dimensions." . format ( partial_batch_dist ) ) ) ) runtime_assertions . append ( assert_util . assert_equal ( expected_partial_batch_shape , partial_batch_dist . batch_shape_tensor ( ) , message = ( "Broadcasting is not supported; " "unexpected batch shape." ) , name = "assert_batch_shape_same" ) ) with tf . control_dependencies ( runtime_assertions ) : return full_sample_and_batch_shape [ : num_broadcast_dims ]
Augment a sample shape to broadcast batch dimensions.

  Computes an augmented sample shape, so that any batch dimensions not
  part of the distribution `partial_batch_dist` are treated as identical
  distributions.

  # partial_batch_dist.batch_shape  = [      7]
  # full_sample_and_batch_shape     = [3, 4, 7]
  # => return an augmented sample shape of [3, 4] so that
  #    partial_batch_dist.sample(augmented_sample_shape) has combined
  #    sample and batch shape of [3, 4, 7].

  Args:
    partial_batch_dist: `tfd.Distribution` instance with batch shape a
      prefix of `full_sample_and_batch_shape`.
    full_sample_and_batch_shape: a Tensor or Tensor-like shape.
    validate_args: if True, check for shape errors at runtime.
  Returns:
    augmented_sample_shape: sample shape such that
      `partial_batch_dist.sample(augmented_sample_shape)` has combined
      sample and batch shape of `full_sample_and_batch_shape`.

  Raises:
    ValueError: if `partial_batch_dist.batch_shape` has more dimensions than
      `full_sample_and_batch_shape`.
    NotImplementedError: if broadcasting would be required to make
      `partial_batch_dist.batch_shape` into a prefix of
      `full_sample_and_batch_shape` .
Augment a sample shape to broadcast batch dimensions .