def machine_config(num_gpus=1, use_tpu=False, master_type=None):
  """Return dict specifying machine config for trainingInput."""
  scale_tier = 'BASIC_GPU'
  if use_tpu:
    scale_tier = 'BASIC_TPU'
  elif num_gpus <= 0:
    scale_tier = 'BASIC'
  elif num_gpus > 1:
    scale_tier = 'CUSTOM'

  config = {'scaleTier': scale_tier}

  if scale_tier == 'CUSTOM':
    assert num_gpus > 1
    if num_gpus not in [4, 8]:
      raise ValueError('Must use exactly 1, 4, or 8 GPUs.')
    config['masterType'] = ('complex_model_m_gpu'
                            if num_gpus == 4 else 'complex_model_l_gpu')

  if master_type:
    config['masterType'] = master_type

  return config