tfp.substrates.jax.bijectors.pack_sequence_as

Returns a Bijector variant of tf.nest.pack_sequence_as.

View aliases

Main aliases

tfp.experimental.substrates.jax.bijectors.pack_sequence_as

tfp.substrates.jax.bijectors.pack_sequence_as(
    example, name='restructure'
)

See also the Restructure bijector for general rearrangements.

Args
example	A Tensor or (potentially nested) collection of Tensors.
name	An optional Python string, inserted into names of TF ops created by this bijector.

Returns
pack	A Bijector whose forward method packs a list of Tensors of the right length into a structure parallel to example, and whose inverse method flattens structures parallel to example into a list of Tensors.

Example

x = tf.constant(1)
example = collections.OrderedDict([
    ('a', [x, x, x]),
    ('b', x)])
bij = tfb.pack_sequence_as(example)
bij.forward([1, 2, 3, 4.])

# Returns
# collections.OrderedDict([
#     ('a', [1, 2, 3]),
#     ('b', 4.)])