DistributionLambda in conjunction with TransformedDistribution and ScalarFunctionWithInferredInverse raises Exception

[MArpogaus]

Hello,

I am using:

• TFP Version 0.16.0
• TF Version 2.8.1

I am training a simple model on a density estimation task.

When using a DistributionLambda in conjunction with a TransformedDistribution and a ScalarFunctionWithInferredInverse bijector i get the following Exception:

ValueError: Exception encountered when calling layer "distribution_lambda" (type DistributionLambda).

Cannot infer tensor rank of objective values.

Call arguments received:
  • inputs=tf.Tensor(shape=(None, 12), dtype=float32)
  • args=<class 'inspect._empty'>
  • kwargs={'training': 'None'}
Error: Process completed with exit code 1.

[Full Traceback]

However, using the same distribution as a custom Loss function works.

[hjlgood]

I would like to use TransformedDistribution as an output of DistributionLambda as well. It would be helpful if there's a way to use more customized distribution other than Normal, Binomial etc.