tfp.substrates.jax.util.TransformedVariable

Variable tracking object which applies a bijector upon convert_to_tensor.

Inherits From: DeferredTensor

View aliases

Main aliases

tfp.experimental.substrates.jax.util.TransformedVariable

tfp.substrates.jax.util.TransformedVariable(
    initial_value, bijector, dtype=None, name=None, **kwargs
)

Example

from tensorflow_probability.python.internal.backend.jax.compat import v2 as tf
import tensorflow_probability as tfp; tfp = tfp.substrates.jax
tfb = tfp.bijectors

positive_variable = tfp.util.TransformedVariable(1., bijector=tfb.Exp())

positive_variable
# ==> <TransformedVariable: dtype=float32, shape=[], fn=exp>

# Note that the initial value corresponds to the transformed output.
tf.convert_to_tensor(positive_variable)
# ==> 1.

positive_variable.pretransformed_input
# ==> <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.0>

# Operators work with `TransformedVariable`.
positive_variable + 1.
# ==> 2.

# It is also possible to assign values to a TransformedVariable
with tf.control_dependencies([positive_variable.assign_add(2.)]):
  positive_variable
# ==> 3.

A common use case for the `TransformedVariable` is to fit constrained
parameters. E.g.:

```python
from tensorflow_probability.python.internal.backend.jax.compat import v2 as tf
import tensorflow_probability as tfp; tfp = tfp.substrates.jax
tfb = tfp.bijectors
tfd = tfp.distributions

trainable_normal = tfd.Normal(
    loc=tf.Variable(0.),
    scale=tfp.util.TransformedVariable(1., bijector=tfb.Exp()))

trainable_normal.loc
# ==> <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.0>

trainable_normal.scale
# ==> <TransformedVariable: dtype=float32, shape=[], fn=exp>

with tf.GradientTape() as tape:
  negloglik = -trainable_normal.log_prob(0.5)
g = tape.gradient(negloglik, trainable_normal.trainable_variables)
# ==> (-0.5, 0.75)

opt = tf.optimizers.Adam(learning_rate=0.05)
loss = tf.function(lambda: -trainable_normal.log_prob(0.5))
for _ in range(int(1e3)):
  opt.minimize(loss, trainable_normal.trainable_variables)
trainable_normal.mean()
# ==> 0.5
trainable_normal.stddev()
# ==> (approximately) 0.0075

Args
`initial_value`	A `Tensor`, or Python object convertible to a `Tensor`, which is the initial value for the `TransformedVariable`. The underlying untransformed `tf.Variable` will be initialized with `bijector.inverse(initial_value)`. Can also be a callable with no argument that returns the initial value when called.
`bijector`	A `Bijector`-like instance which defines the transformations applied to the underlying `tf.Variable`.
`dtype`	`tf.dtype.DType` instance or otherwise valid `dtype` value to `tf.convert_to_tensor(..., dtype)`. Default value: `None` (i.e., `bijector.dtype`).
`name`	Python `str` representing the underlying `tf.Variable`'s name. Default value: `None`.
`**kwargs`	Keyword arguments forward to `tf.Variable`.

Attributes
`also_track`	Additional variables tracked by tf.Module in self.trainable_variables.
`bijector`
`dtype`	Represents the type of the elements in a `Tensor`.
`initializer`	The initializer operation for the underlying variable.
`name`	The string name of this object.
`pretransformed_input`	Input to `transform_fn`.
`shape`	Represents the shape of a `Tensor`.
`trainable_variables`
`transform_fn`	Function which characterizes the `Tensor`ization of this object.
`variables`

Methods

`assign`

View source

assign(
    value, **_
)

`assign_add`

View source

assign_add(
    value, **_
)

`assign_sub`

View source

assign_sub(
    value, **_
)

`numpy`

View source

numpy()

Returns (copy of) deferred values as a NumPy array or scalar.

`set_shape`

View source

set_shape(
    shape
)

Updates the shape of this pretransformed_input.

This method can be called multiple times, and will merge the given shape with the current shape of this object. It can be used to provide additional information about the shape of this object that cannot be inferred from the graph alone.