Can't replicated Masked Autoregressive Flow paper results

I'm trying to train MAFs but I haven't been able to replicate the results reported in the paper using the tensorflow probability implementation.

# Generate data -- as in Figure 1 in [Papamakarios et al. (2017)][2]).
n = 2000
x2 = np.random.randn(n).astype(dtype=np.float32) * 2.
x1 = np.random.randn(n).astype(dtype=np.float32) + (x2 * x2 / 4.)
data = np.stack([x1, x2], axis=-1)

# Density estimation with MADE.
made = tfb.AutoregressiveNetwork(params=2, hidden_units=[512, 512], activation='relu')
biject = tfb.Chain([tfb.MaskedAutoregressiveFlow(made)] * 8)
distribution = tfd.TransformedDistribution(
    distribution=tfd.Normal(loc=0., scale=1.),
    bijector=biject,
    event_shape=[2])

# Construct and fit model.  
x_ = tfkl.Input(shape=(2,), dtype=tf.float32)  
log_prob_ = distribution.log_prob(x_)  
model = tfk.Model(x_, log_prob_)

model.compile(optimizer=tf.optimizers.Adam(),  
              loss=lambda _, log_prob: -log_prob)`  
batch_size = 25  
model.fit(x=data,  
          y=np.zeros((n, 0), dtype=np.float32),  
          batch_size=batch_size,  
          epochs=100000,  
          shuffle=True,  
          verbose=1)

# Use the fitted distribution.  
distribution.sample((2,))  
distribution.log_prob(np.ones((3, 2), dtype=np.float32))

Using TF 2.1 And TFP 0.9.0
I've also tried without chaining bijectors

The model will train and just get stuck in a minima right away and hasn't come close to what the paper reports or other examples online.

Start by running the supplied training script with TensorFlow 2.1 and TFP 0.9.0, then compare the chained and unchained MaskedAutoregressiveFlow cases with the paper's result. No repository file or test is named, so locate the MaskedAutoregressiveFlow and AutoregressiveNetwork entry points first. Done means the reproduction gap is explained or corrected and the expected behavior is verified.