Merge pull request #2272 from ncfrey/tf-type-fix

  assert scores[classification_metric.name] > .9

# def test_normalizing_flow_model_reload():

#   """Test that RobustMultitaskRegressor can be reloaded correctly."""

#   from deepchem.models.normalizing_flows import NormalizingFlow, NormalizingFlowModel

#   import tensorflow_probability as tfp

#   tfd = tfp.distributions

#   tfb = tfp.bijectors

#   tfk = tf.keras

#   tfk.backend.set_floatx('float64')

def test_normalizing_flow_model_reload():

  """Test that NormalizingFlowModel can be reloaded correctly."""

  from deepchem.models.normalizing_flows import NormalizingFlow, NormalizingFlowModel

  import tensorflow_probability as tfp

  tfd = tfp.distributions

  tfb = tfp.bijectors

  tfk = tf.keras

#   model_dir = tempfile.mkdtemp()

  model_dir = tempfile.mkdtemp()

#   Made = tfb.AutoregressiveNetwork(

#       params=2, hidden_units=[512, 512], activation='relu')

  Made = tfb.AutoregressiveNetwork(

      params=2, hidden_units=[512, 512], activation='relu', dtype='float64')

#   flow_layers = [tfb.MaskedAutoregressiveFlow(shift_and_log_scale_fn=Made)]

#   # 3D Multivariate Gaussian base distribution

#   nf = NormalizingFlow(

#       base_distribution=tfd.MultivariateNormalDiag(

#           loc=np.zeros(2), scale_diag=np.ones(2)),

#       flow_layers=flow_layers)

  flow_layers = [tfb.MaskedAutoregressiveFlow(shift_and_log_scale_fn=Made)]

  # 3D Multivariate Gaussian base distribution

  nf = NormalizingFlow(

      base_distribution=tfd.MultivariateNormalDiag(

          loc=np.zeros(2), scale_diag=np.ones(2)),

      flow_layers=flow_layers)

#   nfm = NormalizingFlowModel(nf, model_dir=model_dir)

  nfm = NormalizingFlowModel(nf, model_dir=model_dir)

#   target_distribution = tfd.MultivariateNormalDiag(loc=np.array([1., 0.]))

#   dataset = dc.data.NumpyDataset(X=target_distribution.sample(96))

#   final = nfm.fit(dataset, nb_epoch=1)

  target_distribution = tfd.MultivariateNormalDiag(loc=np.array([1., 0.]))

  dataset = dc.data.NumpyDataset(X=target_distribution.sample(96))

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

#   x = np.zeros(2)

#   lp1 = nfm.flow.log_prob(x).numpy()

  x = np.zeros(2)

  lp1 = nfm.flow.log_prob(x).numpy()

#   assert nfm.flow.sample().numpy().shape == (2,)

  assert nfm.flow.sample().numpy().shape == (2,)

#   reloaded_model = NormalizingFlowModel(nf, model_dir=model_dir)

#   reloaded_model.restore()

  reloaded_model = NormalizingFlowModel(nf, model_dir=model_dir)

  reloaded_model.restore()

#   # Check that reloaded model can sample from the distribution

#   assert reloaded_model.flow.sample().numpy().shape == (2,)

  # Check that reloaded model can sample from the distribution

  assert reloaded_model.flow.sample().numpy().shape == (2,)

#   lp2 = reloaded_model.flow.log_prob(x).numpy()

  lp2 = reloaded_model.flow.log_prob(x).numpy()

#   # Check that density estimation is same for reloaded model

#   assert np.all(lp1 == lp2)

  # Check that density estimation is same for reloaded model

  assert np.all(lp1 == lp2)

def test_robust_multitask_regressor_reload():