def model_fn(features, labels, mode,  # pylint: disable=unused-argument
               params):
    """Wrapper function to select model mode.

    This function is called by tf.estimator.train_and_evaluate function in the
    background.

    Args:
      features: features
      labels: unused labels
      mode: tf.estimator.ModeKeys.PREDICT or TRAIN or EVAL
      params: extra params

    Returns:
      A tf.estimator spec

    """
    emql = EmQL(name, params, data_loader)

    # Define behaviors for different operationstrain / eval / pred
    if mode == tf_estimator.ModeKeys.TRAIN:
      loss, tensors = emql.get_tf_model(name, features, params)
      optimizer = tf.train.AdamOptimizer(params['learning_rate'])
      gvs = optimizer.compute_gradients(loss)
      capped_gvs = [(tf.clip_by_norm(grad, params['gradient_clip']), var)
                    for grad, var in gvs if grad is not None]
      train_op = optimizer.apply_gradients(
          capped_gvs, global_step=tf.train.get_global_step())
      return tf_estimator.EstimatorSpec(
          mode=mode, train_op=train_op, loss=loss)

    elif mode == tf_estimator.ModeKeys.EVAL:
      loss, tensors = emql.get_tf_model(name, features, params)

      if name == 'mixture':
        evaluations = dict()
        evaluations.update(emql.run_tf_evaluation(
            tensors['membership_logits'], tensors['membership_labels'],
            prefix='membership_'))
        evaluations.update(emql.run_tf_evaluation(
            tensors['intersection_logits'], tensors['intersection_labels'],
            prefix='intersection_'))
        evaluations.update(emql.run_tf_evaluation(
            tensors['union_logits'], tensors['union_labels'],
            prefix='union_'))
        evaluations.update(emql.run_tf_evaluation(
            tensors['set_follows_logits'], tensors['set_follows_labels'],
            prefix='set_follows_'))
      else:
        evaluations = emql.run_tf_evaluation(
            tensors['logits'], tensors['labels'])
      return tf_estimator.EstimatorSpec(
          mode=mode, loss=loss, eval_metric_ops=evaluations)

    elif mode == tf_estimator.ModeKeys.PREDICT:
      loss, tensors = emql.get_tf_model(name, features, params)
      predictions = emql.get_tf_prediction(name, features, tensors)
      return tf_estimator.EstimatorSpec(mode=mode, predictions=predictions)

    else:
      raise ValueError('illegal mode %r' % mode)
Wrapper function to select model mode.

    This function is called by tf.estimator.train_and_evaluate function in the
    background.

    Args:
      features: features
      labels: unused labels
      mode: tf.estimator.ModeKeys.PREDICT or TRAIN or EVAL
      params: extra params

    Returns:
      A tf.estimator spec

    
language/emql/model.py