tfp.experimental.substrates.jax.bijectors.Chain

View source on GitHub

Bijector which applies a sequence of bijectors.

Inherits From: Bijector

Example Use:

chain = Chain([Exp(), Softplus()], name="one_plus_exp")

Results in:

  • Forward:
 exp = Exp()
 softplus = Softplus()
 Chain([exp, softplus]).forward(x)
 = exp.forward(softplus.forward(x))
 = tf.exp(tf.log(1. + tf.exp(x)))
 = 1. + tf.exp(x)
 ```

* Inverse:

 ```python
 exp = Exp()
 softplus = Softplus()
 Chain([exp, softplus]).inverse(y)
 = softplus.inverse(exp.inverse(y))
 = tf.log(tf.exp(tf.log(y)) - 1.)
 = tf.log(y - 1.)
 ```

<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>

<tr>
<td>
`bijectors`
</td>
<td>
Python `list` of bijector instances. An empty list makes this
bijector equivalent to the `Identity` bijector.
</td>
</tr><tr>
<td>
`validate_args`
</td>
<td>
Python `bool` indicating whether arguments should be
checked for correctness.
</td>
</tr><tr>
<td>
`parameters`
</td>
<td>
Locals dict captured by subclass constructor, to be used for
copy/slice re-instantiation operators.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
Python `str`, name given to ops managed by this object. Default:
E.g., `Chain([Exp(), Softplus()]).name == "chain_of_exp_of_softplus"`.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2"><h2 class="add-link">Raises</h2></th></tr>

<tr>
<td>
`ValueError`
</td>
<td>
if bijectors have different dtypes.
</td>
</tr>
</table>





<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>

<tr>
<td>
`bijectors`
</td>
<td>

</td>
</tr><tr>
<td>
`dtype`
</td>
<td>
dtype of `Tensor`s transformable by this distribution.
</td>
</tr><tr>
<td>
`forward_min_event_ndims`
</td>
<td>
Returns the minimal number of dimensions bijector.forward operates on.
</td>
</tr><tr>
<td>
`graph_parents`
</td>
<td>
Returns this `Bijector`'s graph_parents as a Python list.
</td>
</tr><tr>
<td>
`inverse_min_event_ndims`
</td>
<td>
Returns the minimal number of dimensions bijector.inverse operates on.
</td>
</tr><tr>
<td>
`is_constant_jacobian`
</td>
<td>
Returns true iff the Jacobian matrix is not a function of x.

Note: Jacobian matrix is either constant for both forward and inverse or
neither.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
Returns the string name of this `Bijector`.
</td>
</tr><tr>
<td>
`parameters`
</td>
<td>
Dictionary of parameters used to instantiate this `Bijector`.
</td>
</tr><tr>
<td>
`trainable_variables`
</td>
<td>

</td>
</tr><tr>
<td>
`validate_args`
</td>
<td>
Returns True if Tensor arguments will be validated.
</td>
</tr><tr>
<td>
`variables`
</td>
<td>

</td>
</tr>
</table>



## Methods

<h3 id="forward"><code>forward</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L1001-L1017">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>forward(
    x, name='forward', **kwargs
)
</code></pre>

Returns the forward `Bijector` evaluation, i.e., X = g(Y).


<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`x`
</td>
<td>
`Tensor`. The input to the 'forward' evaluation.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
The name to give this op.
</td>
</tr><tr>
<td>
`**kwargs`
</td>
<td>
Named arguments forwarded to subclass implementation.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>
<tr class="alt">
<td colspan="2">
`Tensor`.
</td>
</tr>

</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Raises</th></tr>

<tr>
<td>
`TypeError`
</td>
<td>
if `self.dtype` is specified and `x.dtype` is not
`self.dtype`.
</td>
</tr><tr>
<td>
`NotImplementedError`
</td>
<td>
if `_forward` is not implemented.
</td>
</tr>
</table>



<h3 id="forward_dtype"><code>forward_dtype</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L1379-L1395">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>forward_dtype(
    dtype, name='forward_dtype', **kwargs
)
</code></pre>

Returns the dtype of the output of the forward transformation.


<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`dtype`
</td>
<td>
`tf.dtype`, or nested structure of `tf.dtype`s, of the input to
`forward`.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
The name to give this op.
</td>
</tr><tr>
<td>
`**kwargs`
</td>
<td>
Named arguments forwarded to subclass implementation.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>
<tr class="alt">
<td colspan="2">
`tf.dtype` or nested structure of `tf.dtype`s of the output of `forward`.
</td>
</tr>

</table>



<h3 id="forward_event_shape"><code>forward_event_shape</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L913-L927">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>forward_event_shape(
    input_shape
)
</code></pre>

Shape of a single sample from a single batch as a `TensorShape`.

Same meaning as `forward_event_shape_tensor`. May be only partially defined.

<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`input_shape`
</td>
<td>
`TensorShape` indicating event-portion shape passed into
`forward` function.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>

<tr>
<td>
`forward_event_shape_tensor`
</td>
<td>
`TensorShape` indicating event-portion shape
after applying `forward`. Possibly unknown.
</td>
</tr>
</table>



<h3 id="forward_event_shape_tensor"><code>forward_event_shape_tensor</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L886-L906">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>forward_event_shape_tensor(
    input_shape, name='forward_event_shape_tensor'
)
</code></pre>

Shape of a single sample from a single batch as an `int32` 1D `Tensor`.


<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`input_shape`
</td>
<td>
`Tensor`, `int32` vector indicating event-portion shape
passed into `forward` function.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
name to give to the op
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>

<tr>
<td>
`forward_event_shape_tensor`
</td>
<td>
`Tensor`, `int32` vector indicating
event-portion shape after applying `forward`.
</td>
</tr>
</table>



<h3 id="forward_log_det_jacobian"><code>forward_log_det_jacobian</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L1339-L1367">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>forward_log_det_jacobian(
    x, event_ndims, name='forward_log_det_jacobian', **kwargs
)
</code></pre>

Returns both the forward_log_det_jacobian.


<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`x`
</td>
<td>
`Tensor`. The input to the 'forward' Jacobian determinant evaluation.
</td>
</tr><tr>
<td>
`event_ndims`
</td>
<td>
Number of dimensions in the probabilistic events being
transformed. Must be greater than or equal to
`self.forward_min_event_ndims`. The result is summed over the final
dimensions to produce a scalar Jacobian determinant for each event, i.e.
it has shape `rank(x) - event_ndims` dimensions.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
The name to give this op.
</td>
</tr><tr>
<td>
`**kwargs`
</td>
<td>
Named arguments forwarded to subclass implementation.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>
<tr class="alt">
<td colspan="2">
`Tensor`, if this bijector is injective.
If not injective this is not implemented.
</td>
</tr>

</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Raises</th></tr>

<tr>
<td>
`TypeError`
</td>
<td>
if `self.dtype` is specified and `y.dtype` is not
`self.dtype`.
</td>
</tr><tr>
<td>
`NotImplementedError`
</td>
<td>
if neither `_forward_log_det_jacobian`
nor {`_inverse`, `_inverse_log_det_jacobian`} are implemented, or
this is a non-injective bijector.
</td>
</tr>
</table>



<h3 id="inverse"><code>inverse</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L1069-L1087">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>inverse(
    y, name='inverse', **kwargs
)
</code></pre>

Returns the inverse `Bijector` evaluation, i.e., X = g^{-1}(Y).


<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`y`
</td>
<td>
`Tensor`. The input to the 'inverse' evaluation.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
The name to give this op.
</td>
</tr><tr>
<td>
`**kwargs`
</td>
<td>
Named arguments forwarded to subclass implementation.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>
<tr class="alt">
<td colspan="2">
`Tensor`, if this bijector is injective.
If not injective, returns the k-tuple containing the unique
`k` points `(x1, ..., xk)` such that `g(xi) = y`.
</td>
</tr>

</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Raises</th></tr>

<tr>
<td>
`TypeError`
</td>
<td>
if `self.dtype` is specified and `y.dtype` is not
`self.dtype`.
</td>
</tr><tr>
<td>
`NotImplementedError`
</td>
<td>
if `_inverse` is not implemented.
</td>
</tr>
</table>



<h3 id="inverse_dtype"><code>inverse_dtype</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L1397-L1413">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>inverse_dtype(
    dtype, name='inverse_dtype', **kwargs
)
</code></pre>

Returns the dtype of the output of the inverse transformation.


<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`dtype`
</td>
<td>
`tf.dtype`, or nested structure of `tf.dtype`s, of the input to
`inverse`.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
The name to give this op.
</td>
</tr><tr>
<td>
`**kwargs`
</td>
<td>
Named arguments forwarded to subclass implementation.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>
<tr class="alt">
<td colspan="2">
`tf.dtype` or nested structure of `tf.dtype`s of the output of `inverse`.
</td>
</tr>

</table>



<h3 id="inverse_event_shape"><code>inverse_event_shape</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L961-L975">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>inverse_event_shape(
    output_shape
)
</code></pre>

Shape of a single sample from a single batch as a `TensorShape`.

Same meaning as `inverse_event_shape_tensor`. May be only partially defined.

<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`output_shape`
</td>
<td>
`TensorShape` indicating event-portion shape passed into
`inverse` function.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>

<tr>
<td>
`inverse_event_shape_tensor`
</td>
<td>
`TensorShape` indicating event-portion shape
after applying `inverse`. Possibly unknown.
</td>
</tr>
</table>



<h3 id="inverse_event_shape_tensor"><code>inverse_event_shape_tensor</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L934-L954">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>inverse_event_shape_tensor(
    output_shape, name='inverse_event_shape_tensor'
)
</code></pre>

Shape of a single sample from a single batch as an `int32` 1D `Tensor`.


<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`output_shape`
</td>
<td>
`Tensor`, `int32` vector indicating event-portion shape
passed into `inverse` function.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
name to give to the op
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>

<tr>
<td>
`inverse_event_shape_tensor`
</td>
<td>
`Tensor`, `int32` vector indicating
event-portion shape after applying `inverse`.
</td>
</tr>
</table>



<h3 id="inverse_log_det_jacobian"><code>inverse_log_det_jacobian</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L1264-L1297">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>inverse_log_det_jacobian(
    y, event_ndims, name='inverse_log_det_jacobian', **kwargs
)
</code></pre>

Returns the (log o det o Jacobian o inverse)(y).

Mathematically, returns: `log(det(dX/dY))(Y)`. (Recall that: `X=g^{-1}(Y)`.)

Note that `forward_log_det_jacobian` is the negative of this function,
evaluated at `g^{-1}(y)`.

<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`y`
</td>
<td>
`Tensor`. The input to the 'inverse' Jacobian determinant evaluation.
</td>
</tr><tr>
<td>
`event_ndims`
</td>
<td>
Number of dimensions in the probabilistic events being
transformed. Must be greater than or equal to
`self.inverse_min_event_ndims`. The result is summed over the final
dimensions to produce a scalar Jacobian determinant for each event, i.e.
it has shape `rank(y) - event_ndims` dimensions.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
The name to give this op.
</td>
</tr><tr>
<td>
`**kwargs`
</td>
<td>
Named arguments forwarded to subclass implementation.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>

<tr>
<td>
`ildj`
</td>
<td>
`Tensor`, if this bijector is injective.
If not injective, returns the tuple of local log det
Jacobians, `log(det(Dg_i^{-1}(y)))`, where `g_i` is the restriction
of `g` to the `ith` partition `Di`.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Raises</th></tr>

<tr>
<td>
`TypeError`
</td>
<td>
if `self.dtype` is specified and `y.dtype` is not
`self.dtype`.
</td>
</tr><tr>
<td>
`NotImplementedError`
</td>
<td>
if `_inverse_log_det_jacobian` is not implemented.
</td>
</tr>
</table>



<h3 id="__call__"><code>__call__</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/v0.10.0/tensorflow_probability/python/bijectors/_jax/bijector.py#L794-L879">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>__call__(
    value, name=None, **kwargs
)
</code></pre>

Applies or composes the `Bijector`, depending on input type.

This is a convenience function which applies the `Bijector` instance in
three different ways, depending on the input:

1. If the input is a `tfd.Distribution` instance, return
   `tfd.TransformedDistribution(distribution=input, bijector=self)`.
2. If the input is a `tfb.Bijector` instance, return
   `tfb.Chain([self, input])`.
3. Otherwise, return `self.forward(input)`

<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`value`
</td>
<td>
A `tfd.Distribution`, `tfb.Bijector`, or a `Tensor`.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
Python `str` name given to ops created by this function.
</td>
</tr><tr>
<td>
`**kwargs`
</td>
<td>
Additional keyword arguments passed into the created
`tfd.TransformedDistribution`, `tfb.Bijector`, or `self.forward`.
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>

<tr>
<td>
`composition`
</td>
<td>
A `tfd.TransformedDistribution` if the input was a
`tfd.Distribution`, a `tfb.Chain` if the input was a `tfb.Bijector`, or
a `Tensor` computed by `self.forward`.
</td>
</tr>
</table>


#### Examples

```python
sigmoid = tfb.Reciprocal()(
    tfb.AffineScalar(shift=1.)(
      tfb.Exp()(
        tfb.AffineScalar(scale=-1.))))
# ==> `tfb.Chain([
#         tfb.Reciprocal(),
#         tfb.AffineScalar(shift=1.),
#         tfb.Exp(),
#         tfb.AffineScalar(scale=-1.),
#      ])`  # ie, `tfb.Sigmoid()`

log_normal = tfb.Exp()(tfd.Normal(0, 1))
# ==> `tfd.TransformedDistribution(tfd.Normal(0, 1), tfb.Exp())`

tfb.Exp()([-1., 0., 1.])
# ==> tf.exp([-1., 0., 1.])