Probability: Add bijectors: FillDiagonal and FillScaleDiagonal

[fdtomasi]

While I'd guess the use for it it would be minimal, I found myself implementing a FillDiagonal and a FillScaleDiagonal bijectors, in particular to deal with the variational_inducing_observations_scale parameter of tfd.VariationalGaussianProcess, that reads as

variational_inducing_observations_scale: `float` `Tensor`; the scale
        matrix of the (full-rank Gaussian) variational posterior over function
        values at the inducing points, conditional on observed data. Shape has
        the form `[b1, ..., bB, e2, e2]`, where `b1, ..., bB` is broadcast
        compatible with other parameters and `e2` is the number of inducing
        points.

I like the general framework of VariationalGaussianProcess, even though it would be good to be able to specify a variational_inducing_observations_scale_diag parameter to avoid materialising a scale matrix while only interested in the diagonal. Side note: I briefly tried to look into the code of VariationalGaussianProcess but I was not sure how to deal with variational_gaussian_process._solve_cholesky_factored_system(cholesky_factor, rhs, name=None) when rhs is a vector representing a diagonal matrix.

While I guess this is probably a workaround to deal with the VariationalGaussianProcess case, it is probably useful for other use cases, hence if it makes sense I can submit a PR with the FillDiagonal and FillScaleDiagonal bijectors.
The bijectors are modifications of FillTriangular and FillScaleTril.

class FillDiagonal(bijector.Bijector):
    """Transforms vectors to diagonal matrices.

  Given input with shape `batch_shape + [d]`, produces output with
  shape `batch_shape + [d, d]`.
  [...]

python b = tfb.FillDiagonal() b.forward([1, 2, 3]) # ==> [[1, 0, 0], # [0, 2, 0], # [0, 0, 3]]

and

class FillScaleDiagonal(chain.Chain):
    """Transforms unconstrained vectors to Diagonal matrices with positive diagonal.
  This is implemented as a simple `tfb.Chain` of `tfb.FillDiagonal`
  followed by `tfb.TransformDiagonal`, and provided mostly as a
  convenience. The default setup is somewhat opinionated, using a
  Softplus transformation followed by a small shift (`1e-5`) which
  attempts to avoid numerical issues from zeros on the diagonal.
  [...]
  ```

```python
  b = tfb.FillScaleDiagonal(
       diag_bijector=tfb.Exp(),
       diag_shift=None)
  b.forward(x=[0., 0.])
  # Result: [[1., 0.],
  #          [0., 1.]]
  b.inverse(y=[[1., 0],
               [0, 2]])
  # Result: [log(1), log(2)]

[brianwa84]

+1 vote from me. I ended up using FillScaleTriL, but I agree a diagonal
option would have been nice.

On Sat, Feb 15, 2020, 6:17 AM Federico Tomasi [e-mail]
wrote:

While I'd guess the use for it it would be minimal, I found myself
implementing a FillDiagonal and a FillScaleDiagonal bijectors, in
particular to deal with the variational_inducing_observations_scale
parameter of tfd.VariationalGaussianProcess, that reads as

variational_inducing_observations_scale: float Tensor; the scale
matrix of the (full-rank Gaussian) variational posterior over function
values at the inducing points, conditional on observed data. Shape has
the form [b1, ..., bB, e2, e2], where b1, ..., bB is broadcast
compatible with other parameters and e2 is the number of inducing
points.

I like the general framework of VariationalGaussianProcess, even though
it would be good to be able to specify a
variational_inducing_observations_scale_diag parameter to avoid
materialising a scale matrix while only interested in the diagonal. Side
note: I briefly tried to look into the code of VariationalGaussianProcess
but I was not sure how to deal with variational_gaussian_process._solve_cholesky_factored_system(cholesky_factor,
rhs, name=None) when rhs is a vector representing a diagonal matrix.

While I guess this is probably a workaround to deal with the
VariationalGaussianProcess case, it is probably useful for other use
cases, hence if it makes sense I can submit a PR with the FillDiagonal and
FillScaleDiagonal bijectors.
The bijectors are modifications of FillTriangular and FillScaleTril.

class FillDiagonal(bijector.Bijector):
"""Transforms vectors to diagonal matrices. Given input with shape batch_shape + [d], produces output with shape batch_shape + [d, d]. [...]

b = tfb.FillDiagonal()
b.forward([1, 2, 3])
# ==> [[1, 0, 0],
# [0, 2, 0],
# [0, 0, 3]]

and

class FillScaleDiagonal(chain.Chain):
"""Transforms unconstrained vectors to Diagonal matrices with positive diagonal. This is implemented as a simple tfb.Chain of tfb.FillDiagonal followed by tfb.TransformDiagonal, and provided mostly as a convenience. The default setup is somewhat opinionated, using a Softplus transformation followed by a small shift (1e-5) which attempts to avoid numerical issues from zeros on the diagonal. [...]

b = tfb.FillScaleDiagonal(
diag_bijector=tfb.Exp(),
diag_shift=None)
b.forward(x=[0., 0.])
# Result: [[1., 0.],
# [0., 1.]]
b.inverse(y=[[1., 0],
[0, 2]])
# Result: [log(1), log(2)]

—
You are receiving this because you are subscribed to this thread.
Reply to this email directly, view it on GitHub
https://github.com/tensorflow/probability/issues/791?email_source=notifications&email_token=AFJFSIZDCHA5FUHZXH6QHV3RC7FOHA5CNFSM4KVXYR22YY3PNVWWK3TUL52HS4DFUVEXG43VMWVGG33NNVSW45C7NFSM4INYNOJA,
or unsubscribe
https://github.com/notifications/unsubscribe-auth/AFJFSI3OUK7MFE2F3XOFFLTRC7FOHANCNFSM4KVXYR2Q
.

[srvasude]

Hi,
I'm wondering if an alternative solution would be to accept for variational_inducing_observations_scale, a LinearOperator much like MVNLinearOperator. This would allow for structured lower triangular operators (such as a diagonal matrix), without sacrificing the speed / memory advantages of the structured operation (in this case the diagonal solve and matmul).

[srvasude]

@csuter

[csuter]

Accepting LinOp in VGP SGTM!