def METHOD_NAME(unused_argv):
  ##############################################################################
  # Prepare training and testing data
  ##############################################################################
  (x_train, y_train), (
      x_test,
      y_test), input_shape = keras_test_utils.get_preprocessed_mnist_data()
  ##############################################################################
  # Train and convert a model with 2x2 block config. There's no kernel in tflite
  # supporting this block configuration, so the sparse tensor is densified and
  # the model falls back to dense execution.
  ##############################################################################
  pruning_params = {
      'pruning_schedule':
          ConstantSparsity(FLAGS.sparsity, begin_step=0, frequency=100),
      'block_size': (2, 2)
  }
  model = build_layerwise_model(input_shape, **pruning_params)
  model = train(model, x_train, y_train, x_test, y_test)
  converter = tf.lite.TFLiteConverter.from_keras_model(model)
  # Get a dense model as baseline
  tflite_model_dense = converter.convert()
  tflite_model_path = '/tmp/dense_mnist.tflite'
  with open(tflite_model_path, 'wb') as f:
    f.write(tflite_model_dense)
  # The _experimental_sparsify_model is to enable sparse tensor encoding,
  # otherwise the model is converted as dense.
  converter._experimental_sparsify_model = True
  tflite_model = converter.convert()
  # Check the model is compressed
  print('Compression ratio: ', len(tflite_model) / len(tflite_model_dense))
  tflite_model_path = '/tmp/sparse_mnist_%s_2x2.tflite' % FLAGS.sparsity
  with open(tflite_model_path, 'wb') as f:
    f.write(tflite_model)
  print('evaluate 2x2 model')
  print(keras_test_utils.eval_mnist_tflite(model_content=tflite_model))
  ##############################################################################
  # Train and convert a model with 1x4 block config. There's kernel in tflite
  # with this block configuration, so the model can take advantage of sparse
  # execution and see inference speed-up.
  ##############################################################################
  pruning_params = {
      'pruning_schedule':
          ConstantSparsity(FLAGS.sparsity, begin_step=0, frequency=100),
      # TFLite transposes the weight during conversion, so we need to specify
      # the block as (4, 1) in the training API.
      'block_size': (4, 1)
  }
  model = build_layerwise_model(input_shape, **pruning_params)
  model = train(model, x_train, y_train, x_test, y_test)
  converter = tf.lite.TFLiteConverter.from_keras_model(model)
  converter._experimental_sparsify_model = True
  tflite_model = converter.convert()
  # Check the model is compressed
  print('Compression ratio: ', len(tflite_model) / len(tflite_model_dense))
  tflite_model_path = '/tmp/sparse_mnist_%s_1x4.tflite' % FLAGS.sparsity
  with open(tflite_model_path, 'wb') as f:
    f.write(tflite_model)
  print('evaluate 1x4 model')
  print(keras_test_utils.eval_mnist_tflite(model_content=tflite_model))