def METHOD_NAME(model_dir, params):
  """Construct either an Estimator or TPUEstimator for NCF.
  Args:
    model_dir: The model directory for the estimator
    params: The params dict for the estimator
  Returns:
    An Estimator or TPUEstimator.
  """
  distribution = ncf_common.get_v1_distribution_strategy(params)
  run_config = tf.estimator.RunConfig(train_distribute=distribution,
                                      eval_distribute=distribution)
  model_fn = neumf_model.neumf_model_fn
  if params["use_xla_for_gpu"]:
    # TODO(seemuch): remove the contrib imput
    from tensorflow.contrib.compiler import xla
    logging.info("Using XLA for GPU for training and evaluation.")
    model_fn = xla.estimator_model_fn(model_fn)
  estimator = tf.estimator.Estimator(model_fn=model_fn, model_dir=model_dir,
                                     config=run_config, params=params)
  return estimator