def METHOD_NAME(compile_model=False):
  """Returns an instance of `tf.keras.Model` for use with the MNIST example.
  This code is based on the following target, which unfortunately cannot be
  imported as it is a Python binary, not a library:
  https://github.com/tensorflow/models/blob/master/official/r1/mnist/mnist.py
  Args:
    compile_model: If True, compile the model with a basic optimizer and loss.
  Returns:
    A `tf.keras.Model`.
  """
  # TODO: b/120157713 - Find a way to import this code.
  data_format = 'channels_last'
  input_shape = [28, 28, 1]
  initializer = _DeterministicInitializer(
      tf.keras.initializers.RandomNormal, base_seed=0
  )
  max_pool = tf.keras.layers.MaxPooling2D(
      (2, 2), (2, 2), padding='same', data_format=data_format
  )
  model = tf.keras.Sequential([
      tf.keras.layers.Reshape(target_shape=input_shape, input_shape=(28 * 28,)),
      tf.keras.layers.Conv2D(
          32,
          5,
          padding='same',
          data_format=data_format,
          activation=tf.nn.relu,
          kernel_initializer=initializer(),
      ),
      max_pool,
      tf.keras.layers.Conv2D(
          64,
          5,
          padding='same',
          data_format=data_format,
          activation=tf.nn.relu,
          kernel_initializer=initializer(),
      ),
      max_pool,
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(
          1024, activation=tf.nn.relu, kernel_initializer=initializer()
      ),
      tf.keras.layers.Dropout(0.4, seed=1),
      tf.keras.layers.Dense(10, kernel_initializer=initializer()),
  ])
  if compile_model:
    model.compile(
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        optimizer=tf.keras.optimizers.SGD(learning_rate=0.1),
    )
  return model