description: Builds a trainable blockwise tf.linalg.LinearOperator.
View source on GitHub |
Builds a trainable blockwise tf.linalg.LinearOperator
.
tfp.experimental.vi.util.build_trainable_linear_operator_block(
operators, block_dims=None, batch_shape=(), dtype=None, seed=None, name=None
)
This function returns a trainable blockwise LinearOperator
. If operators
is a flat list, it is interpreted as blocks along the diagonal of the
structure and an instance of tf.linalg.LinearOperatorBlockDiag
is returned.
If operators
is a doubly nested list, then a
tf.linalg.LinearOperatorBlockLowerTriangular
instance is returned, with
the block in row i
column j
(i >= j
) given by operators[i][j]
.
The operators
list may contain LinearOperator
instances, LinearOperator
subclasses, or callables that return LinearOperator
instances. The
dimensions of the blocks are given by block_dims
; this argument may be
omitted if operators
contains only LinearOperator
instances.
# Build a 5x5 trainable `LinearOperatorBlockDiag` given `LinearOperator`
# subclasses and `block_dims`.
op = build_trainable_linear_operator_block(
operators=(tf.linalg.LinearOperatorDiag,
tf.linalg.LinearOperatorLowerTriangular),
block_dims=[3, 2],
dtype=tf.float32)
# Build an 8x8 `LinearOperatorBlockLowerTriangular`, with a callable that
# returns a `LinearOperator` in the upper left block, and `LinearOperator`
# subclasses in the lower two blocks.
op = build_trainable_linear_operator_block(
operators=(
(lambda shape, dtype: tf.linalg.LinearOperatorScaledIdentity(
num_rows=shape[-1], multiplier=tf.Variable(1., dtype=dtype))),
(tf.linalg.LinearOperatorFullMatrix,
tf.linalg.LinearOperatorLowerTriangular))
block_dims=[4, 4],
dtype=tf.float64)
# Build a 6x6 `LinearOperatorBlockDiag` with batch shape `(4,)`. Since
# `operators` contains only `LinearOperator` instances, the `block_dims`
# argument is not necessary.
op = build_trainable_linear_operator_block(
operators=(tf.linalg.LinearOperatorDiag(tf.Variable(tf.ones((4, 3)))),
tf.linalg.LinearOperatorFullMatrix([4.]),
tf.linalg.LinearOperatorIdentity(2)))
Args | |
---|---|
operators
|
A list or tuple containing LinearOperator subclasses,
LinearOperator instances, or callables returning LinearOperator
instances. If the list is flat, a tf.linalg.LinearOperatorBlockDiag
instance is returned. Otherwise, the list must be singly nested, with the
first element of length 1, second element of length 2, etc.; the
elements of the outer list are interpreted as rows of a lower-triangular
block structure, and a tf.linalg.LinearOperatorBlockLowerTriangular
instance is returned. Callables contained in the lists must take three
arguments -- shape , the shape of the tf.Variable instantiating the
LinearOperator , dtype , the tf.dtype of the LinearOperator , and
seed , a seed for generating random values.
|
block_dims
|
List or tuple of integers, representing the sizes of the blocks
along one dimension of the (square) blockwise LinearOperator . If
operators contains only LinearOperator instances, block_dims may be
None and the dimensions are inferred.
|
batch_shape
|
Batch shape of the LinearOperator .
|
dtype
|
tf.dtype of the LinearOperator .
|
seed
|
PRNG seed; see tfp.random.sanitize_seed for details.
|
name
|
str, name for tf.name_scope .
|
Returns | |
---|---|
Trainable instance of tf.linalg.LinearOperatorBlockDiag or
tf.linalg.LinearOperatorBlockLowerTriangular .
|