//  [python]
def model_fn(self, features, labels, mode, params):
    """TPUEstimator compatible model function."""
    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    labelled_features = labels[0]
    labels = tf.cast(labels[1], tf.float32)
    data_shape = features.get_shape().as_list()[1:]
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      z_mean, z_logvar = self.gaussian_encoder(
          features, is_training=is_training)
      z_mean_labelled, _ = self.gaussian_encoder(
          labelled_features, is_training=is_training)

    supervised_loss = []
    mine_ops = []

    for l in range(labels.get_shape().as_list()[1]):
      for r in range(z_mean.get_shape().as_list()[1]):
        label_for_mi = tf.layers.flatten(labels[:, l])
        representation_for_mi = tf.layers.flatten(z_mean_labelled[:, r])
        mi_lr, op_lr = mine(representation_for_mi, label_for_mi,
                            "estimator_network_%d_%d" % (l, r))
        if l != r:
          supervised_loss = supervised_loss + [tf.math.square(mi_lr)]
        mine_ops = mine_ops + [op_lr]
    supervised_loss = tf.reshape(tf.add_n(supervised_loss), [])
    z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar)
    reconstructions = self.decode(z_sampled, data_shape, is_training)
    per_sample_loss = losses.make_reconstruction_loss(features, reconstructions)
    reconstruction_loss = tf.reduce_mean(per_sample_loss)
    kl_loss = compute_gaussian_kl(z_mean, z_logvar)
    standard_vae_loss = tf.add(
        reconstruction_loss, self.beta * kl_loss, name="VAE_loss")
    gamma_annealed = make_annealer(self.gamma_sup, tf.train.get_global_step())
    s2_mine_vae_loss = tf.add(
        standard_vae_loss, gamma_annealed * supervised_loss,
        name="s2_factor_VAE_loss")
    if mode == tf.estimator.ModeKeys.TRAIN:
      optimizer_vae = optimizers.make_vae_optimizer()
      all_variables = tf.trainable_variables()
      encoder_vars = [var for var in all_variables if "encoder" in var.name]
      decoder_vars = [var for var in all_variables if "decoder" in var.name]

      update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
      train_op_vae = optimizer_vae.minimize(
          loss=s2_mine_vae_loss,
          global_step=tf.train.get_global_step(),
          var_list=encoder_vars + decoder_vars)
      train_op = tf.group(train_op_vae, mine_ops, update_ops)
      tf.summary.scalar("reconstruction_loss", reconstruction_loss)
      logging_hook = tf.train.LoggingTensorHook({
          "loss": s2_mine_vae_loss,
          "reconstruction_loss": reconstruction_loss,
          "supervised_loss": supervised_loss,
      },
                                                every_n_iter=50)
      return TPUEstimatorSpec(
          mode=mode,
          loss=s2_mine_vae_loss,
          train_op=train_op,
          training_hooks=[logging_hook])
    elif mode == tf.estimator.ModeKeys.EVAL:
      return TPUEstimatorSpec(
          mode=mode,
          loss=s2_mine_vae_loss,
          eval_metrics=(make_metric_fn("reconstruction_loss", "supervised_loss",
                                       "kl_loss"),
                        [reconstruction_loss, supervised_loss, kl_loss]))
    else:
      raise NotImplementedError("Eval mode not supported.")