tensorflow_probability/python/distributions/vector_diffeomixture.py
def determine_batch_event_shapes(grid, endpoint_affine):
  """Helper to infer batch_shape and event_shape."""
  with tf.name_scope("determine_batch_event_shapes"):
    # grid  # shape: [B, k, q]
    # endpoint_affine     # len=k, shape: [B, d, d]
    batch_shape = grid.shape[:-2]
    batch_shape_tensor = tf.shape(input=grid)[:-2]
    event_shape = None
    event_shape_tensor = None

    def _set_event_shape(shape, shape_tensor):
      if event_shape is None:
        return shape, shape_tensor
      return (tf.broadcast_static_shape(event_shape, shape),
              tf.broadcast_dynamic_shape(event_shape_tensor, shape_tensor))

    for aff in endpoint_affine:
      if aff.shift is not None:
        batch_shape = tf.broadcast_static_shape(batch_shape,
                                                aff.shift.shape[:-1])
        batch_shape_tensor = tf.broadcast_dynamic_shape(
            batch_shape_tensor,
            tf.shape(input=aff.shift)[:-1])
        event_shape, event_shape_tensor = _set_event_shape(
            aff.shift.shape[-1:],
            tf.shape(input=aff.shift)[-1:])

      if aff.scale is not None:
        batch_shape = tf.broadcast_static_shape(batch_shape,
                                                aff.scale.batch_shape)
        batch_shape_tensor = tf.broadcast_dynamic_shape(
            batch_shape_tensor, aff.scale.batch_shape_tensor())
        event_shape, event_shape_tensor = _set_event_shape(
            tf.TensorShape([aff.scale.range_dimension]),
            aff.scale.range_dimension_tensor()[tf.newaxis])

    return batch_shape, batch_shape_tensor, event_shape, event_shape_tensor
def determine_batch_event_shapes(grid, endpoint_affine):
  with tf.name_scope("determine_batch_event_shapes"):
    # grid  # shape: [B, k, q]
    # endpoint_affine     # len=k, shape: [B, d, d]
    batch_shape = grid.shape[:-2]
    batch_shape_tensor = tf.shape(input=grid)[:-2]
    event_shape = None
    event_shape_tensor = None

    def _set_event_shape(shape, shape_tensor):
      if event_shape is None:
        return shape, shape_tensor
      return (tf.broadcast_static_shape(event_shape, shape),
              tf.broadcast_dynamic_shape(event_shape_tensor, shape_tensor))

    for aff in endpoint_affine:
      if aff.shift is not None:
        batch_shape = tf.broadcast_static_shape(batch_shape,
                                                aff.shift.shape[:-1])
        batch_shape_tensor = tf.broadcast_dynamic_shape(
            batch_shape_tensor,
            tf.shape(input=aff.shift)[:-1])
        event_shape, event_shape_tensor = _set_event_shape(
            aff.shift.shape[-1:],
            tf.shape(input=aff.shift)[-1:])

      if aff.scale is not None:
        batch_shape = tf.broadcast_static_shape(batch_shape,
                                                aff.scale.batch_shape)
        batch_shape_tensor = tf.broadcast_dynamic_shape(
            batch_shape_tensor, aff.scale.batch_shape_tensor())
        event_shape, event_shape_tensor = _set_event_shape(
            tf.TensorShape([aff.scale.range_dimension]),
            aff.scale.range_dimension_tensor()[tf.newaxis])

    return batch_shape, batch_shape_tensor, event_shape, event_shape_tensor
def determine_batch_event_shapes ( grid , endpoint_affine ) : with tf . name_scope ( "determine_batch_event_shapes" ) : # grid  # shape: [B, k, q] # endpoint_affine     # len=k, shape: [B, d, d] batch_shape = grid . shape [ : - 2 ] batch_shape_tensor = tf . shape ( input = grid ) [ : - 2 ] event_shape = None event_shape_tensor = None def _set_event_shape ( shape , shape_tensor ) : if event_shape is None : return shape , shape_tensor return ( tf . broadcast_static_shape ( event_shape , shape ) , tf . broadcast_dynamic_shape ( event_shape_tensor , shape_tensor ) ) for aff in endpoint_affine : if aff . shift is not None : batch_shape = tf . broadcast_static_shape ( batch_shape , aff . shift . shape [ : - 1 ] ) batch_shape_tensor = tf . broadcast_dynamic_shape ( batch_shape_tensor , tf . shape ( input = aff . shift ) [ : - 1 ] ) event_shape , event_shape_tensor = _set_event_shape ( aff . shift . shape [ - 1 : ] , tf . shape ( input = aff . shift ) [ - 1 : ] ) if aff . scale is not None : batch_shape = tf . broadcast_static_shape ( batch_shape , aff . scale . batch_shape ) batch_shape_tensor = tf . broadcast_dynamic_shape ( batch_shape_tensor , aff . scale . batch_shape_tensor ( ) ) event_shape , event_shape_tensor = _set_event_shape ( tf . TensorShape ( [ aff . scale . range_dimension ] ) , aff . scale . range_dimension_tensor ( ) [ tf . newaxis ] ) return batch_shape , batch_shape_tensor , event_shape , event_shape_tensor
Helper to infer batch_shape and event_shape.
Helper to infer batch_shape and event_shape .