Refactor tests

    """

    return Lambda(lambda x: -tf.reduce_mean(self.flow.log_prob(x + 1e-10, training=True)))(output)

    return Lambda(lambda x: -tf.reduce_mean(self.flow.log_prob(x, training=True)))(output)

class NormalizingFlowLayer(object):

tfb = tfp.bijectors

class TestNormalizingFlow(unittest.TestCase):

  def setUp(self):

def test_normalizing_flow():

  flow_layers = [

      tfb.RealNVP(

              hidden_layers=[8, 8]))

  ]

  # 3D Multivariate Gaussian base distribution

    self.nf = NormalizingFlow(

  nf = NormalizingFlow(

      base_distribution=tfd.MultivariateNormalDiag(loc=[0., 0., 0.]),

      flow_layers=flow_layers)

    self.nfm = NormalizingFlowModel(self.nf, batch_size=1)

  nfm = NormalizingFlowModel(nf)

  # Must be float32 for RealNVP

    self.dataset = NumpyDataset(

  dataset = NumpyDataset(

      X=np.random.rand(5, 3).astype(np.float32),

      y=np.random.rand(5,),

      ids=np.arange(5))

  def test_simple_flow(self):

    """Tests a simple flow of one RealNVP layer."""

  # Tests a simple flow of one RealNVP layer.

    X = self.nfm.flow.sample()

  X = nfm.flow.sample()

  x1 = tf.zeros([3])

    x2 = self.dataset.X[0]

  x2 = dataset.X[0]

  # log likelihoods should be negative

    assert self.nfm.flow.log_prob(X).numpy() < 0

    assert self.nfm.flow.log_prob(x1).numpy() < 0

    assert self.nfm.flow.log_prob(x2).numpy() < 0

  assert nfm.flow.log_prob(X).numpy() < 0

  assert nfm.flow.log_prob(x1).numpy() < 0

  assert nfm.flow.log_prob(x2).numpy() < 0

  # # Fit model

    final = self.nfm.fit(self.dataset, nb_epoch=5)

  final = nfm.fit(dataset, nb_epoch=5)

  assert final > 0