The equivalent of LinearOperator in JAX

[pawel-czyz]

This issue is related to TensorFlow Probability on JAX.

The tutorials of multivariate Student and multivariate normal use tf.linalg.LinearOperatorLowerTriangular.

However, this object does not exist in JAX. I wonder what is the correct way to build these distributions?
Passing a JAX array leads to an error:

TypeError: Expected argument `scale` to be instance of `LinearOperator`. Found: `[[ 0.6  0.   0. ]
 [ 0.2  0.5  0. ]
 [ 0.1 -0.3  0.4]]`.

Any help would be much appreciated!

Additional notes

• Issue a bug when using MultivariateStudentTLinearOperator in DistributionLamda during training #1595 may be related with this one.

[brianwa84]

If you've done import tensorflow_probability.substrates.jax as tfp, then you can find linear operators in: jtf = tfp.tf2jax; dir(jtf.linalg)

[pawel-czyz]

It resolved the issue, thank you very much! 🙂