def _forward_event_shape_tensor(self, input_shape, is_inverse=False):
ndims = ps.size(input_shape)
indices = ps.reshape(ps.add(self.axis, ndims), shape=[-1, 1])
extra_sizes = ps.reduce_sum(self.paddings, axis=-1)
update_fn = (ps.tensor_scatter_nd_sub if is_inverse else
ps.tensor_scatter_nd_add)
return update_fn(ps.identity(input_shape), indices, extra_sizes)
def batch_shape_tensor(self, name='batch_shape_tensor'):
"""Shape of a single sample from a single event index as a 1-D `Tensor`.
The batch dimensions are indexes into independent, non-identical
parameterizations of this distribution.
Args:
name: name to give to the op
Returns:
batch_shape: `Tensor`.
"""
with tf.name_scope(name):
if tensorshape_util.is_fully_defined(self.batch_shape):
batch_shape = self.batch_shape
else:
batch_shape = ps.shape(self._log_prob_increment_fn())
return ps.identity(
ps.convert_to_shape_tensor(batch_shape, name='batch_shape'))