tfp.edward2.tape(
*args,
**kwds
)
Context manager for recording interceptable executions onto a tape.
Similar to tf.GradientTape
, operations are recorded if they are executed
within this context manager. In addition, the operation must be registered
(wrapped) as ed.interceptable
.
Yields:
tape
: OrderedDict where operations are recorded in sequence. Keys are thename
keyword argument to the operation (typically, a random variable'sname
) and values are the corresponding output of the operation. If the operation has no name, it is not recorded.
Examples
from tensorflow_probability import edward2 as ed
def probabilistic_matrix_factorization():
users = ed.Normal(0., 1., sample_shape=[5000, 128], name="users")
items = ed.Normal(0., 1., sample_shape=[7500, 128], name="items")
ratings = ed.Normal(loc=tf.matmul(users, items, transpose_b=True),
scale=0.1,
name="ratings")
return ratings
with ed.tape() as model_tape:
ratings = probabilistic_matrix_factorization()
assert model_tape["users"].shape == (5000, 128)
assert model_tape["items"].shape == (7500, 128)
assert model_tape["ratings"] == ratings