text: # correctness test
  # MNIST model taken from github/tensorflow/models
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  import os
  import sys
  
  import tensorflow as tf
  from tensorflow.examples.tutorials.mnist import input_data
  
  # import memory_saving_gradients from ..
  module_path=os.path.dirname(os.path.abspath(__file__))
  sys.path.append(module_path+'/..')
  import memory_saving_gradients
  import mem_util
  
  TEST_DEVICE='/cpu:0'
  USE_REAL_DATA = False
  FLAGS_data_dir='/tmp/mnist_data'
  FLAGS_model_dir='/tmp/mnist_model'
  FLAGS_batch_size=1
  FLAGS_data_format=None
  
  # from tensorflow.contrib.learn.python.learn.datasets.mnist import read_data_sets
  
    
  def train_dataset(data_dir):
    """Returns a tf.data.Dataset yielding (image, label) pairs for training."""
    data = input_data.read_data_sets(data_dir, one_hot=True).train
    return tf.data.Dataset.from_tensor_slices((data.images, data.labels))
  
  
  def mnist_model(inputs, mode, data_format):
    """Takes the MNIST inputs and mode and outputs a tensor of logits."""
    # Input Layer
    # Reshape X to 4-D tensor: [batch_size, width, height, channels]
    # MNIST images are 28x28 pixels, and have one color channel
    inputs = tf.reshape(inputs, [-1, 28, 28, 1])
  
    if data_format is None:
      # When running on GPU, transpose the data from channels_last (NHWC) to
      # channels_first (NCHW) to improve performance.
      # See https://www.tensorflow.org/performance/performance_guide#data_formats
      data_format = ('channels_first'
                     if tf.test.is_gpu_available() else 'channels_last')
  
    if data_format == 'channels_first':
      inputs = tf.transpose(inputs, [0, 3, 1, 2])
  
    # Convolutional Layer #1
    # Computes 32 features using a 5x5 filter with ReLU activation.
    # Padding is added to preserve width and height.
    # Input Tensor Shape: [batch_size, 28, 28, 1]
    # Output Tensor Shape: [batch_size, 28, 28, 32]
    conv1 = tf.layers.conv2d(
        inputs=inputs,
        filters=32,
        kernel_size=[5, 5],
        padding='same',
        activation=tf.nn.relu,
        data_format=data_format)
  
    # Pooling Layer #1
    # First max pooling layer with a 2x2 filter and stride of 2
    # Input Tensor Shape: [batch_size, 28, 28, 32]
    # Output Tensor Shape: [batch_size, 14, 14, 32]
    pool1 = tf.layers.max_pooling2d(
        inputs=conv1, pool_size=[2, 2], strides=2, data_format=data_format)
  
    # Convolutional Layer #2
    # Computes 64 features using a 5x5 filter.
    # Padding is added to preserve width and height.
    # Input Tensor Shape: [batch_size, 14, 14, 32]
    # Output Tensor Shape: [batch_size, 14, 14, 64]
    conv2 = tf.layers.conv2d(
        inputs=pool1,
        filters=64,
        kernel_size=[5, 5],
        padding='same',
        activation=tf.nn.relu,
        data_format=data_format)
  
    # Pooling Layer #2
    # Second max pooling layer with a 2x2 filter and stride of 2
    # Input Tensor Shape: [batch_size, 14, 14, 64]
    # Output Tensor Shape: [batch_size, 7, 7, 64]
    pool2 = tf.layers.max_pooling2d(
        inputs=conv2, pool_size=[2, 2], strides=2, data_format=data_format)
  
    # Flatten tensor into a batch of vectors
    # Input Tensor Shape: [batch_size, 7, 7, 64]
    # Output Tensor Shape: [batch_size, 7 * 7 * 64]
    pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
  
    # Dense Layer
    # Densely connected layer with 1024 neurons
    # Input Tensor Shape: [batch_size, 7 * 7 * 64]
    # Output Tensor Shape: [batch_size, 1024]
    dense = tf.layers.dense(inputs=pool2_flat, units=1024, activation=tf.nn.relu)
  
    # Add dropout operation; 0.6 probability that element will be kept
    dropout = tf.layers.dropout(
        inputs=dense, rate=0.4, training=(mode == tf.estimator.ModeKeys.TRAIN))
  
    # Logits layer
    # Input Tensor Shape: [batch_size, 1024]
    # Output Tensor Shape: [batch_size, 10]
    logits = tf.layers.dense(inputs=dropout, units=10)
    return logits
  
  GLOBAL_PROFILE = True
  DUMP_TIMELINES = False
  run_metadata = True
  def sessrun(*args, **kwargs):
    global sess, run_metadata
    
    if not GLOBAL_PROFILE:
      return sess.run(*args, **kwargs)
    
    run_metadata = tf.RunMetadata()
  
    kwargs['options'] = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    kwargs['run_metadata'] = run_metadata
    result = sess.run(*args, **kwargs)
    first_entry = args[0]
    if isinstance(first_entry, list):
      if len(first_entry) == 0 and len(args) == 1:
        return None
      first_entry = first_entry[0]
  
    if DUMP_TIMELINES:
      name = first_entry.name
      name = name.replace('/', '-')
  
      tl = timeline.Timeline(run_metadata.step_stats)
      ctf = tl.generate_chrome_trace_format()
      with open('timelines/%s.json'%(name,), 'w') as f:
        f.write(ctf)
      with open('timelines/%s.pbtxt'%(name,), 'w') as f:
        f.write(str(run_metadata))
  
    return result
  
  
  def create_session():
    from tensorflow.core.protobuf import rewriter_config_pb2
    optimizer_options = tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0)
    config = tf.ConfigProto(operation_timeout_in_ms=150000, graph_options=tf.GraphOptions(optimizer_options=optimizer_options))
    config.graph_options.rewrite_options.constant_folding = rewriter_config_pb2.RewriterConfig.OFF
    config.graph_options.place_pruned_graph = True
    return tf.Session(config=config)
  
  
  sess = None
  def train_mnist():
    global sess
  
    # restrict to cpu:0
    tf.reset_default_graph()
    tf.set_random_seed(1)
    np.random.seed(1)
    tf_dev = tf.device(TEST_DEVICE)
    tf_dev.__enter__()
  
    #  FLAGS = parse_flags()
    # Train the model
  
  
    # replace Dataset ops with constant images because gradient rewriting
    # tries to differentiate graphs containing IteratorGetNext
    # TODO: make it work with Dataset ops
    images = tf.Variable(tf.random_uniform((FLAGS_batch_size, 28**2)))
    labels = tf.Variable(tf.concat([tf.ones((FLAGS_batch_size, 1)),
                                    tf.zeros((FLAGS_batch_size, 9))], axis=1))
    def train_input_fn():
      dataset = train_dataset(FLAGS_data_dir)
      dataset = dataset.batch(FLAGS_batch_size)
      (images, labels) = dataset.make_one_shot_iterator().get_next()
      num_images = FLAGS_batch_size
      return (images[:num_images], labels[:num_images])
  
    if USE_REAL_DATA:
      images, labels = train_input_fn()
  #    images = tf.stop_gradient(images)
  #    labels = tf.stop_gradient(labels)
  
    
    logits = mnist_model(images, tf.estimator.ModeKeys.TRAIN, 'channels_last')
    cross_entropy = tf.losses.softmax_cross_entropy(logits=logits,
                                                    onehot_labels=labels)
    loss = cross_entropy
    optimizer = tf.train.GradientDescentOptimizer(learning_rate=1e-2)
  
    vars = tf.trainable_variables()
    grads = tf.gradients(loss, vars)
    grads_and_vars = zip(grads, vars)
    train_op = optimizer.apply_gradients(grads_and_vars)
    
    sess = create_session()
    sess.run(tf.global_variables_initializer())
    print("Loss %.5f" %(sess.run(loss)))
    for i in range(10):
      sessrun(train_op)
      mem_use = mem_util.peak_memory(run_metadata)[TEST_DEVICE]/1e6
      print("Loss %.5f, memory %.2f MB" %(sess.run(loss), mem_use))
  
    # should print something like this for actual dataset
    # 2.12764
    # 1.87759
    # 1.54445
    # 1.29149
    # 1.18474
    # 0.884424
    # 0.69454
    # 0.770236
    # 0.629259
    # 0.654465
  
    assert sess.run(loss) < 100
  
  
  def test_correctness(capsys):
    # enable printing during successful test run under pytest, uncomment these
    # if capsys:
    #   pytest_decorator = capsys.disabled()
    #   pytest_decorator.__enter__()
  
    # Loss 0.01803, memory 399.10 MB
    # Loss 0.00002, memory 399.10 MB
    # Loss 0.00000, memory 399.10 MB
    # Running with memory saving
    # Extracting /tmp/mnist_data/train-images-idx3-ubyte.gz
    # Extracting /tmp/mnist_data/train-labels-idx1-ubyte.gz
    # Extracting /tmp/mnist_data/t10k-images-idx3-ubyte.gz
    # Extracting /tmp/mnist_data/t10k-labels-idx1-ubyte.gz
    # Loss 0.07283, memory 380.72 MB
    # Loss 0.00398, memory 351.23 MB
    # Loss 0.00035, memory 351.23 MB
  
    def grads(ys, xs, grad_ys=None, **kwargs):
      return memory_saving_gradients.gradients(ys, xs, grad_ys,
                                               checkpoints='speed', **kwargs)
    old_grads = tf.gradients
    tf.__dict__["gradients"] = grads
    print("Running with memory saving")
    
    train_mnist()
  
    print("\nRunning with regular gradient")
    tf.__dict__["gradients"] = old_grads
    train_mnist()
  
  def main(unused_argv):
    test_correctness(None)
  
  if __name__ == '__main__':
    main([])
  