tensorflow_probability/python/distributions/batch_reshape.py
def _call_and_reshape_output(
      self,
      fn,
      event_shape_list=None,
      static_event_shape_list=None,
      extra_kwargs=None):
    """Calls `fn` and appropriately reshapes its output."""
    # Note: we take `extra_kwargs` as a dict rather than `**extra_kwargs`
    # because it is possible the user provided extra kwargs would itself
    # have `fn`, `event_shape_list`, `static_event_shape_list` and/or
    # `extra_kwargs` as keys.
    with tf.control_dependencies(self._runtime_assertions):
      if event_shape_list is None:
        event_shape_list = [self._event_shape_tensor()]
      if static_event_shape_list is None:
        static_event_shape_list = [self.event_shape]
      new_shape = tf.concat(
          [self._batch_shape_unexpanded] + event_shape_list, axis=0)
      result = tf.reshape(fn(**extra_kwargs) if extra_kwargs else fn(),
                          new_shape)
      if (tensorshape_util.rank(self.batch_shape) is not None and
          tensorshape_util.rank(self.event_shape) is not None):
        event_shape = tf.TensorShape([])
        for rss in static_event_shape_list:
          event_shape = tensorshape_util.concatenate(event_shape, rss)
        static_shape = tensorshape_util.concatenate(
            self.batch_shape, event_shape)
        tensorshape_util.set_shape(result, static_shape)
      return result
def _call_and_reshape_output(
      self,
      fn,
      event_shape_list=None,
      static_event_shape_list=None,
      extra_kwargs=None):
    """Calls `fn` and appropriately reshapes its output."""
    # Note: we take `extra_kwargs` as a dict rather than `**extra_kwargs`
    # because it is possible the user provided extra kwargs would itself
    # have `fn`, `event_shape_list`, `static_event_shape_list` and/or
    # `extra_kwargs` as keys.
    with tf.control_dependencies(self._runtime_assertions):
      if event_shape_list is None:
        event_shape_list = [self._event_shape_tensor()]
      if static_event_shape_list is None:
        static_event_shape_list = [self.event_shape]
      new_shape = tf.concat(
          [self._batch_shape_unexpanded] + event_shape_list, axis=0)
      result = tf.reshape(fn(**extra_kwargs) if extra_kwargs else fn(),
                          new_shape)
      if (tensorshape_util.rank(self.batch_shape) is not None and
          tensorshape_util.rank(self.event_shape) is not None):
        event_shape = tf.TensorShape([])
        for rss in static_event_shape_list:
          event_shape = tensorshape_util.concatenate(event_shape, rss)
        static_shape = tensorshape_util.concatenate(
            self.batch_shape, event_shape)
        tensorshape_util.set_shape(result, static_shape)
      return result
def _call_and_reshape_output ( self , fn , event_shape_list = None , static_event_shape_list = None , extra_kwargs = None ) : # Note: we take `extra_kwargs` as a dict rather than `**extra_kwargs` # because it is possible the user provided extra kwargs would itself # have `fn`, `event_shape_list`, `static_event_shape_list` and/or # `extra_kwargs` as keys. with tf . control_dependencies ( self . _runtime_assertions ) : if event_shape_list is None : event_shape_list = [ self . _event_shape_tensor ( ) ] if static_event_shape_list is None : static_event_shape_list = [ self . event_shape ] new_shape = tf . concat ( [ self . _batch_shape_unexpanded ] + event_shape_list , axis = 0 ) result = tf . reshape ( fn ( * * extra_kwargs ) if extra_kwargs else fn ( ) , new_shape ) if ( tensorshape_util . rank ( self . batch_shape ) is not None and tensorshape_util . rank ( self . event_shape ) is not None ) : event_shape = tf . TensorShape ( [ ] ) for rss in static_event_shape_list : event_shape = tensorshape_util . concatenate ( event_shape , rss ) static_shape = tensorshape_util . concatenate ( self . batch_shape , event_shape ) tensorshape_util . set_shape ( result , static_shape ) return result
Calls `fn` and appropriately reshapes its output.
Calls fn and appropriately reshapes its output .