Conditional MaskedAutoregressive Flow outputs NAN

Hi, I am experimenting with the bijector (MaskedAutoregressive Flow). Essentially I want it to be a mapping from a length 20 vector drawn from some base distribution to some other distribution, conditioned on some input value (which is a length 10 vector). Below is my code.

import tensorflow_probability as tfp
import tensorflow as tf
import numpy as np
tfd = tfp.distributions
tfb = tfp.bijectors

input = np.ones(shape=(20)).astype(np.float32)
condition = np.random.normal(size=(10)).astype(np.float32) + 10

fn = tfb.AutoregressiveNetwork(params=2, event_shape=20, conditional=True, conditional_event_shape=10, hidden_units=[10, 10])
bijector = tfb.MaskedAutoregressiveFlow(fn)

print(bijector.forward(input, conditional_input=condition))

Result of running the code gives me:

tf.Tensor( [nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan nan], shape=(20,), dtype=float32)

I don't think my input or condition vectors are big enough to cause exploding values. Any ideas on what I might be doing wrong? Thanks!

For reference, I am using tensorflow-probability = 0.11.0

Start by reproducing the supplied Python snippet with tensorflow-probability 0.11.0 and inspect the conditional AutoregressiveNetwork and MaskedAutoregressiveFlow path it exercises. Determine why the forward result contains NaNs, then add or update a regression test for the reported input and condition shapes so the result is finite.