DistributionLambda in conjunction with TransformedDistribution and ScalarFunctionWithInferredInverse raises Exception

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.

Start with cml/old_faithful.py and reproduce the failure using DistributionLambda, TransformedDistribution, and ScalarFunctionWithInferredInverse with the reported TFP and TF versions. Compare this path with the custom loss that works, using the linked full traceback to locate the rank-inference failure; done means the distribution can be used in DistributionLambda without the exception.