tensorflow_probability/python/distributions/batch_reshape.py
def _call_reshape_input_output(self, fn, x, extra_kwargs=None):
    """Calls `fn`, appropriately reshaping its input `x` and 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` and/or `x` as a key.
    with tf.control_dependencies(self._runtime_assertions +
                                 self._validate_sample_arg(x)):
      sample_shape, static_sample_shape = self._sample_shape(x)
      old_shape = tf.concat(
          [
              sample_shape,
              self.distribution.batch_shape_tensor(),
              self.event_shape_tensor(),
          ],
          axis=0)
      x_reshape = tf.reshape(x, old_shape)
      result = fn(x_reshape, **extra_kwargs) if extra_kwargs else fn(x_reshape)
      new_shape = tf.concat(
          [
              sample_shape,
              self._batch_shape_unexpanded,
          ], axis=0)
      result = tf.reshape(result, new_shape)
      if (tensorshape_util.rank(static_sample_shape) is not None and
          tensorshape_util.rank(self.batch_shape) is not None):
        new_shape = tensorshape_util.concatenate(static_sample_shape,
                                                 self.batch_shape)
        tensorshape_util.set_shape(result, new_shape)
      return result
def _call_reshape_input_output(self, fn, x, 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` and/or `x` as a key.
    with tf.control_dependencies(self._runtime_assertions +
                                 self._validate_sample_arg(x)):
      sample_shape, static_sample_shape = self._sample_shape(x)
      old_shape = tf.concat(
          [
              sample_shape,
              self.distribution.batch_shape_tensor(),
              self.event_shape_tensor(),
          ],
          axis=0)
      x_reshape = tf.reshape(x, old_shape)
      result = fn(x_reshape, **extra_kwargs) if extra_kwargs else fn(x_reshape)
      new_shape = tf.concat(
          [
              sample_shape,
              self._batch_shape_unexpanded,
          ], axis=0)
      result = tf.reshape(result, new_shape)
      if (tensorshape_util.rank(static_sample_shape) is not None and
          tensorshape_util.rank(self.batch_shape) is not None):
        new_shape = tensorshape_util.concatenate(static_sample_shape,
                                                 self.batch_shape)
        tensorshape_util.set_shape(result, new_shape)
      return result
def _call_reshape_input_output ( self , fn , x , 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` and/or `x` as a key. with tf . control_dependencies ( self . _runtime_assertions + self . _validate_sample_arg ( x ) ) : sample_shape , static_sample_shape = self . _sample_shape ( x ) old_shape = tf . concat ( [ sample_shape , self . distribution . batch_shape_tensor ( ) , self . event_shape_tensor ( ) , ] , axis = 0 ) x_reshape = tf . reshape ( x , old_shape ) result = fn ( x_reshape , * * extra_kwargs ) if extra_kwargs else fn ( x_reshape ) new_shape = tf . concat ( [ sample_shape , self . _batch_shape_unexpanded , ] , axis = 0 ) result = tf . reshape ( result , new_shape ) if ( tensorshape_util . rank ( static_sample_shape ) is not None and tensorshape_util . rank ( self . batch_shape ) is not None ) : new_shape = tensorshape_util . concatenate ( static_sample_shape , self . batch_shape ) tensorshape_util . set_shape ( result , new_shape ) return result
Calls `fn`, appropriately reshaping its input `x` and output.
Calls fn appropriately reshaping its input x and output .