str-tf-patch.diff

diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index b322779e..a64a3827 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -186,6 +186,7 @@ py_library(
         "//tensorflow/python/ops/parallel_for",
         "//tensorflow/python/ops/ragged",
         "//tensorflow/python/ops/signal",
+        "//tensorflow/python/ops/stropt",
         "//tensorflow/python/profiler",
         "//tensorflow/python/saved_model",
         "//tensorflow/python/tools:module_util",
diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py
index 907b118c..e175f926 100644
--- a/tensorflow/python/keras/engine/training.py
+++ b/tensorflow/python/keras/engine/training.py
@@ -20,6 +20,7 @@ from __future__ import print_function
 
 import collections
 import json
+from re import S
 import numpy as np
 
 from tensorflow.python import tf2
@@ -1006,19 +1007,23 @@ class Model(network.Network):
       outputs = [
           training_v2_utils._non_none_constant_value(v) for v in outputs]  # pylint: disable=protected-access
     else:
+      # logging.info("=======graph mode training batch start=======")
       x = training_utils.ModelInputs(x).as_list()
       ins = x + (y or []) + (sample_weights or [])
 
       if not isinstance(K.symbolic_learning_phase(), int):
         ins += [True]  # Add learning phase value.
-
       self._update_sample_weight_modes(sample_weights=sample_weights)
+      # logging.info("Prepare train function")
       self._make_train_function()
+      # logging.info("Invoke training")
+      # logging.info(f"Computational graph size={len(self._graph.get_operations())}")
       outputs = self.train_function(ins)  # pylint: disable=not-callable
 
     if reset_metrics:
       self.reset_metrics()
 
+    # logging.info("=======graph mode training batch end=======")
     if len(outputs) == 1:
       return outputs[0]
     return outputs
diff --git a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
index 6b916fc7..7ebc97ad 100644
--- a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
+++ b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
@@ -22,8 +22,11 @@ from __future__ import print_function
 
 import abc
 import functools
+from logging import info
 
 import six
+import configparser
+import os
 
 from tensorflow.python.distribute import distribution_strategy_context as distribute_ctx
 from tensorflow.python.distribute import reduce_util as ds_reduce_util
@@ -386,7 +389,12 @@ class OptimizerV2(trackable.Trackable):
     params = nest.flatten(params)
     with backend.get_graph().as_default(), backend.name_scope(self._name +
                                                               "/gradients"):
-      grads = gradients.gradients(loss, params)
+      from tensorflow.python.ops.stropt.str_strategy import str_gradients
+      config = configparser.ConfigParser()
+      config.read(os.environ['HOME'] + "/.str/strategy.conf")
+      grads = str_gradients(loss, params, config=config, optimizer_name=self._name)
+
+      # grads = gradients.gradients(loss, params) # vanilla
       for grad, param in zip(grads, params):
         if grad is None:
           raise ValueError("Variable {} has `None` for gradient. "
diff --git a/tensorflow/python/keras/optimizers.py b/tensorflow/python/keras/optimizers.py
index 10ce8336..bb3d3fa4 100644
--- a/tensorflow/python/keras/optimizers.py
+++ b/tensorflow/python/keras/optimizers.py
@@ -17,6 +17,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
+import logging
 
 import six
 from six.moves import zip  # pylint: disable=redefined-builtin
diff --git a/tensorflow/python/lib/core/bfloat16.cc b/tensorflow/python/lib/core/bfloat16.cc
index 7d170113..fde3a837 100644
--- a/tensorflow/python/lib/core/bfloat16.cc
+++ b/tensorflow/python/lib/core/bfloat16.cc
@@ -490,7 +490,7 @@ bool RegisterBfloat16Cast(int numpy_type, bool cast_is_safe) {
 }
 
 template <typename InType, typename OutType, typename Functor>
-void BinaryUFunc(char** args, npy_intp const* dimensions, npy_intp const* steps,
+void BinaryUFunc(char** args, npy_intp* dimensions, npy_intp* steps,
                  void* data) {
   const char* i0 = args[0];
   const char* i1 = args[1];
@@ -506,7 +506,7 @@ void BinaryUFunc(char** args, npy_intp const* dimensions, npy_intp const* steps,
 }
 
 template <typename Functor>
-void CompareUFunc(char** args, npy_intp const* dimensions, npy_intp const* steps,
+void CompareUFunc(char** args, npy_intp* dimensions, npy_intp* steps,
                   void* data) {
   BinaryUFunc<bfloat16, npy_bool, Functor>(args, dimensions, steps, data);
 }
diff --git a/tensorflow/python/ops/stropt/BUILD b/tensorflow/python/ops/stropt/BUILD
new file mode 100644
index 00000000..774b2767
--- /dev/null
+++ b/tensorflow/python/ops/stropt/BUILD
@@ -0,0 +1,21 @@
+package(
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+py_library(
+    name = "stropt",
+    srcs = [
+        "__init__.py",
+        "memory_saving_gradients.py",
+        "str_strategy.py",
+        "strategy_parser.py",
+        "lms.py",
+        "topos.py",
+    ],
+    srcs_version = "PY2AND3",
+    deps = [
+        "//tensorflow/contrib/graph_editor:graph_editor_py",
+    ]
+)
+
diff --git a/tensorflow/python/ops/stropt/__init__.py b/tensorflow/python/ops/stropt/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tensorflow/python/ops/stropt/lms.py b/tensorflow/python/ops/stropt/lms.py
new file mode 100644
index 00000000..8b8da6c6
--- /dev/null
+++ b/tensorflow/python/ops/stropt/lms.py
@@ -0,0 +1,1040 @@
+# (C) Copyright IBM Corp. 2018. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""LMS
+"""
+import re
+import tensorflow as tf
+import tensorflow.contrib.graph_editor as ge
+from tensorflow.contrib.graph_editor import util
+
+import time
+from six.moves import queue as Queue
+from tensorflow.python.ops.stropt import topos
+from enum import Enum
+
+
+class CTRLD_Strategy(Enum):
+    CHAIN_RULE = 1
+    DIRECT_ORDER = 2
+
+# Operations with these types will be excluded from swapping
+ATOMIC_TYPES = {'Const', 'Mul', 'Add',
+                'Identity', 'Assign', 'VariableV2',
+                'Reshape', 'Shape', 'ShapeN',
+                'Placeholder'}
+
+
+class LMS(object):
+    """LMS class for Large Model Support (LMS).
+
+    The `LMS` object statically modifies a model by swapping its tensors
+    to the host so that the model can be trained with the limited memory
+    of GPUs.
+
+    Tensors those are generated by forward operations and consumed by
+    backward operations are candidates for swapping. The `LMS` object will
+    automatically find these tensors.
+
+    Swapping is done by cutting the link between a forward operation and
+    its backward operation, then replacing the link by inserting `identity`
+    operations on the host. In theory, this procedure does not have any
+    effect on the training convergence as well as inference task.
+    """
+    def __init__(self, optimizer_scopes, graph=None,
+                 starting_scope=None,
+                 starting_op_names=None,
+                 excl_scopes=set(),
+                 incl_scopes=set(),
+                 excl_types=set(),
+                 incl_types=set(),
+                 lb=1, ub=10000,
+                 n_tensors=-1,
+                 fuse_swapins=False,
+                 ctrld_strategy="chain_rule",
+                 swap_branches=False,
+                 branch_threshold=0,
+                 debug=False,
+                 debug_level=1,
+                 cpu_device="/cpu:0",
+                 solved_strategy=None):
+        """Create an LMS object to edit the graph for supporting large model.
+
+        Args:
+          graph: the graph we will modify for LMS. This should be the graph of
+            user-defined neural network.
+          optimizer_scopes: a set of scopes for the optimizers/solvers.
+          starting_scope: tensors that are reachable from the operations in
+            this scope will be swapped for LMS. Set this to the scope of the
+            first layer if we would like to modify the whole graph.
+          starting_op_names: tensors that are reachable from the operations with
+            these names will be swapped for LMS.
+          excl_scopes: a set of scopes for operations whose tensors will not
+            be swapped out to the host. Default `empty`.
+          incl_scopes: a set of scopes for operations whose tensors will be
+            swapped out to the host. Default `empty`.
+          excl_types: a set of types for operations whose tensors will not be
+            swapped out to the host. Default `empty`.
+          incl_types: a set of types for operations whose tensors will be
+            swapped out to the host. Default `empty`.
+          n_tensors: the number of tensors for LMS, counting from the
+            `starting_scope`. To turn off LMS, set `n_tensors` to `0`.
+            Default `-1` (all reachable tensors will be swapped for LMS).
+          lb: lower-bound value for LMS. A tensor will be swapped in during the
+            backward phase at least `lb` nodes before it in the graph.
+            Default `1`.
+          ub: upper-bound value for LMS. Default `10000`.
+          fuse_swapins: Fuse "close" swap-in operations into one operation.
+            This may improve the performance. Default `False`.
+          ctrld_strategy: Two strategies to find control dependency ops for
+            swapin ops: `chain_rule` and `direct_order`. `chain_rule` strategy
+            starts from a forward operation, goes forward and finds a corresponding
+            backward operation to be a control dependency operation. `direct_order`
+            strategy directly gets a backward ops in the topological order to be
+            a control dependency operation. Both strategies depend on `lb` and `ub`
+            to choose a control dependency operation. While the `direct_order` is
+            more exact than `chain_rule` in relation to `lb` and `ub`, it experimentally
+            often results in smaller maximum batch size than `chain_rule`.
+            Default `chain_rule`.
+          swap_branches: If True, LMS will swap tensors in branches in the
+            forward phase. Default `False`.
+          branch_threshold: If `swap_branches` is enabled and the
+            topological-sort distance between the consuming operation and generating
+            operation of a tensor is greater than `branch_threshold`, then swap the
+            tensor. Default `0`.
+          debug: debug mode for LMS. Default `False`.
+          debug_level: Debug level for LMS (1 or 2). Default `1`.
+          cpu_device: the device we would like swap tensors to.
+        """
+        if not optimizer_scopes:
+            raise ValueError('A least one optimizer scope is required.')
+
+        self._graph = graph
+        self._optimizer_scopes = optimizer_scopes
+        self._excl_scopes = excl_scopes
+        self._incl_scopes = incl_scopes
+        self._excl_types = excl_types
+        self._incl_types = incl_types
+        self._starting_scope = starting_scope
+        self._starting_op_names = starting_op_names
+        self._lb = lb  # lowerbound
+        self._ub = ub  # upperbound
+        self._n_tensors = n_tensors
+        self._fuse_swapins = fuse_swapins
+        if ctrld_strategy == "chain_rule":
+            self._ctrld_strategy = CTRLD_Strategy.CHAIN_RULE
+        elif ctrld_strategy == "direct_order":
+            self._ctrld_strategy = CTRLD_Strategy.DIRECT_ORDER
+        else:
+            self._ctrld_strategy = CTRLD_Strategy.CHAIN_RULE
+
+        self._swap_branches = swap_branches
+        self._branch_threshold = branch_threshold
+
+        self._excl_types |= ATOMIC_TYPES
+
+        self._excl_ops = set()
+        self._incl_ops = set()
+        self._grad_ops = set()
+        self._topo_sort = None
+        self._cpu_device = cpu_device
+        self._debug = debug
+        self._debug_level = debug_level
+
+        # keep log of tensors on host
+        self._incpu_count = 0
+
+        # store a dictionary of visited ops to avoid multiple visits
+        self._ops_dict = {}
+
+        # Solved strategy of STR
+        self._ss = solved_strategy
+
+    def _build_gradient_ops(self):
+        """Return a set of operations in the backward phase.
+
+        Operations in the backward phase are determined by its scope.
+        """
+        for scope in self._optimizer_scopes:
+            ops_for_scope = set(ge.filter_ops_from_regex(
+                ge.make_list_of_op(self._graph), "^{}".format(scope)))
+            if not ops_for_scope:
+                self._log_info('No operations were found with optimizer '
+                               'scope {}.'.format(scope))
+            self._grad_ops.update(ops_for_scope)
+        if not self._grad_ops:
+            raise ValueError('No operations were found with optimizer '
+                             'scopes {}.'.format(self._optimizer_scopes))
+
+
+    def _get_seed_ops(self):
+        """Return a list of `tf.Operation` used as a starting point for LMS
+        to traverse the graph.
+
+        If a starting scope is given, the ops in this scope will be used.
+        Otherwise, this method automatically searches for starting ops.
+        """
+        # seep ops for search
+        seed_ops = set()
+        ops = ge.make_list_of_op(self._graph)
+        if self._starting_scope:
+            scope_ops = set(ge.filter_ops_from_regex(ops,
+                                                     "^{}".format(self._starting_scope)))
+            if not scope_ops:
+                raise ValueError('No operations were found in starting '
+                                 'scope {}.'.format(self._starting_scope))
+            seed_ops |= scope_ops
+
+        if self._starting_op_names:
+            for name in self._starting_op_names:
+                name_ops = set(ge.filter_ops_from_regex(ops,
+                                                        "^{}$".format(name)))
+                if not name_ops:
+                    raise ValueError('No starting operation was found with '
+                                     'name {}.'.format(name))
+                seed_ops |= name_ops
+
+        seed_ops = list(seed_ops)
+        if not seed_ops:
+            candidates = set()
+            non_grad_ops = [op
+                            for op in self._graph.get_operations()
+                            if not (op in self._grad_ops)]
+            for op in non_grad_ops:
+                for t in op.outputs:
+                    frontier_ops = set(util.get_consuming_ops(t))
+                    if (frontier_ops & self._grad_ops):
+                        candidates.add(op)
+                        break
+
+            # ordering an operation by how much it covers the other ops
+            tmp_dict = {}
+            max_nelems = -1
+            for op in candidates:
+                nelems = len(set(ge.get_forward_walk_ops(op, within_ops=non_grad_ops,
+                                                         inclusive=False)) &
+                             candidates)
+                if nelems > 0:
+                    tmp_dict[op] = nelems
+                    max_nelems = nelems if (nelems > max_nelems) else max_nelems
+
+            # seed ops will cover most of the forward ops
+            seed_ops = [k for k, v in tmp_dict.items() if v == max_nelems]
+        return seed_ops
+
+    def _filter_scopes_and_types(self, within_ops, scopes, types):
+        """Return ops in within_ops that are in `scopes` or have a type
+        in `types`.
+
+        Args:
+          within_ops: an object convertible to a list of `tf.Operation`.
+          scopes: a list of scope path.
+          types: a list of tf.DataType.
+        Return:
+          A set of `tf.Operation`.
+        """
+        ret_ops = set()
+        for scope in scopes:
+            ops = set(ge.get_name_scope_ops(within_ops, scope))
+            if not ops:
+                raise ValueError('No operations were found with scope'
+                                 ' {}.'.format(scope))
+            ret_ops |= ops
+
+        found_types = set()
+        type_ops = set()
+        for op in within_ops:
+            if op.type in types:
+                found_types.add(op.type)
+                type_ops.add(op)
+
+        # We remove ATOMIC_TYPES from the input list of types because
+        # it is a constant and not user input. We only want to error if a
+        # user provided type is not found.
+        missing_types = types - found_types - ATOMIC_TYPES
+        if missing_types:
+            raise ValueError('No operations were found with types: '
+                             ' {}.'.format(str(missing_types)))
+
+        ret_ops |= type_ops
+        return ret_ops
+
+    def _get_forward_walk_ops(self, op, inclusive=True):
+        """ A wrapper of `tensorflow.contrib.graph_editor.get_forward_walk_ops`
+        """
+        if op in self._ops_dict:
+            if inclusive:
+                return self._ops_dict[op]
+            else:
+                return list(set(self._ops_dict[op]) - {op})
+        else:
+            ret = ge.get_forward_walk_ops(op)
+            self._ops_dict[op] = ret
+            if inclusive:
+                return ret
+            else:
+                return list(set(ret) - {op})
+
+    def run(self, graph=None):
+        """Edit the graph by adding swapin and swapout ops.
+
+        Swapin and swapout ops are in the host.
+
+        The graph is modified in-place.
+
+        Return:
+          a set of added ops.
+        """
+        if graph:
+            self._graph = graph
+
+        if self._n_tensors == 0:
+            self._log_info("LMS is disabled and will not modify the model.")
+            return  # turn off LMS
+        elif self._n_tensors < 0:
+            self._n_tensors = 0  # swap all tensors (default)
+
+        if not self._graph:
+            raise ValueError('The dataflow graph is required but has not been'
+                             ' provided.')
+
+        self._log_info("Editing model for LMS")
+        self._print_configuration()
+        start_time = time.time()
+
+        self._build_gradient_ops()        
+        seed_ops = self._get_seed_ops() # first op of this graph
+
+        self._log_info(
+            "Starting ops: {}".format(
+                [(op.name, op.type) for op in seed_ops]), 1)
+
+        reachable_ops = set()
+        for seed_op in seed_ops:
+            reachable_ops |= set(self._get_forward_walk_ops(seed_op))
+
+        for op in reachable_ops:
+            if 'lms/swap' in op.name:
+                self._log_info('This model has already been updated with LMS '
+                               'swap operations. LMS will not re-process it.')
+                return
+        # exclusive ops
+        self._excl_ops = self._filter_scopes_and_types(reachable_ops,
+                                                       self._excl_scopes,
+                                                       self._excl_types)
+        # inclusive ops
+        self._incl_ops = self._filter_scopes_and_types(reachable_ops,
+                                                       self._incl_scopes,
+                                                       self._incl_types)
+
+        reachable_ops -= self._grad_ops
+
+        # build a topological sort
+        self._topo_sort = topos.TOPOS(seed_ops, self._grad_ops)
+        self._topo_sort.build()
+        for i in range(0, self._topo_sort.size):
+            self._log_info("[{}]: {}".format(
+                i, [op.name for op in self._topo_sort.get_ops(i)]), 1)
+        
+        if self._ss is None:
+            self._do_action(seed_ops)
+        else:
+            # STR instrumentation
+            fwd_ops = set(ge.make_list_of_op(self._graph)).difference(list(self._grad_ops))
+            records = self._ss.get_swap_record(list(fwd_ops), list(self._grad_ops))
+            self._log_info("Total {} records are found from STR strategy.".format(len(records)))
+            self._str_swap_action(records)
+
+        # check the validation of the new model
+        new_reachable_ops = set()
+        for seed_op in seed_ops:
+            new_reachable_ops |= set(ge.get_forward_walk_ops(seed_op))
+        new_reachable_ops -= self._grad_ops
+        if (new_reachable_ops >= reachable_ops):
+            self._log_info("Edited model is valid and logically equivalent to the original one")
+            self._log_info("Added {} ops into the model".format(len(new_reachable_ops - reachable_ops)))
+        else:
+            self._log_info("Edited model is invalid. Running this may produce unexpected result")
+
+        self._log_info("Editing model for LMS, took: {} ms".format(
+            (time.time()-start_time)*1000))
+        self._log_info(
+            "{} tensors will be swapped out(in) to(from) the host".format(
+                self._incpu_count))
+        return (new_reachable_ops - reachable_ops)
+
+    def _str_swap_action(self, records):
+        """ Edit the computing graph with swapping nodes. The editing is based on the solved strategy of STR
+
+        Args:
+            records (List[SwapRecord]): a list of swap records, which contains the swap out/in operations.
+        """        
+        for i, record in enumerate(records):
+            self._log_info("Processing record: {}".format(record))
+            swapout_op = self._add_swapout(record.swap_out_source_ts.op, record.swap_out_source_ts)
+            self._incpu_count = self._incpu_count + 1     
+            ge.add_control_inputs(swapout_op, [record.swap_out_control_op])
+            self._log_info("Add swap out node {}".format([swapout_op]))
+            
+            # create swap_in nodes
+            for control_op, dest_ops in zip(record.swap_in_control_ops, record.swap_in_dest_ops):
+                self._log_info("Start add swap in node for dest. {}, control at {}".format(dest_ops, control_op))
+                # one swap in tensor may be used by multiple nodes
+                swapin_op = self._add_swapin_str(swapout_op, dest_ops, record.swap_out_source_ts)
+                if swapin_op is not None:
+                    ge.add_control_inputs(swapin_op, [control_op])
+                    self._log_info("Add swap in node: {}, linking to dest ops:{}".format([swapin_op], dest_ops))
+        
+
+    def _do_action(self, src_ops):
+        """Add swapin and swapout ops for ops that are reachable from `src_ops`.
+
+        Args:
+          src_ops: a list of `tf.Operation`
+        """
+        open_set = Queue.Queue()
+        closed_set = set()
+
+        for op in src_ops:
+            open_set.put(op)
+
+        while not open_set.empty():
+            src_op = open_set.get()
+
+            # get next ops before the graph is changed
+            next_ops = set()
+            for t in src_op.outputs:
+                frontier_ops = set(util.get_consuming_ops(t))
+                next_ops |= frontier_ops - self._grad_ops
+
+            # do action for src_op
+            self._insert_swap_nodes(src_op)
+            if self._swapped_max_tensors():
+                return
+
+            for op in next_ops:
+                if op in closed_set:
+                    continue
+                if op not in open_set.queue:
+                    open_set.put(op)
+
+            closed_set.add(src_op)
+
+    def _fuse_swapin_ops(self, src_op, swapout_op, bw_frontier_ops, ts0):
+        """Fuse all swapin ops that swaps in the same tensor.
+
+        This method does an in-place modification to the graph.
+
+        Args:
+          src_op: a `tf.Operation`.
+          swapout_op: a `tf.Operation`.
+          bw_frontier_ops: a set of `tf.Operation`.
+          ts0: a `tf.Tensor`.
+
+        Return:
+          A set of `tf.Operation` that cannot be fused.
+        """
+        fuse_bw_frontier_ops = {
+            op for op in bw_frontier_ops
+            if self._topo_sort.get_order(op) > 0}
+        if len(fuse_bw_frontier_ops) >= 2:
+            with tf.device(self._cpu_device):
+                swap_in = tf.identity(ts0, name="lms/swapin")
+
+            # Connect: swap_out -> swap_in
+            self._connect_ops(swapout_op, swap_in.op)
+            self._excl_ops.add(swap_in.op)
+
+            # reuse swap_in tensors
+            for op in fuse_bw_frontier_ops:
+                # Connect: swap_in -> dest
+                input_idx = ge.sgv(
+                    op, graph=self._graph).input_index(ts0)
+                self._connect_ops(swap_in.op, op, remap_inputs=True,
+                                  idx=input_idx)
+
+                self._log_info(
+                    "{} (order {}) reuses tensor {}".format(
+                        op.name,
+                        self._topo_sort.get_order(op),
+                        ts0.name),
+                    1)
+
+            # control dependency -> swap_in
+            min_order = self._topo_sort.size + 1
+            earliest_op = None
+            for op in fuse_bw_frontier_ops:
+                order = self._topo_sort.get_order(op)
+                if order < min_order:
+                    min_order = order
+                    earliest_op = op
+            if earliest_op:
+                self._add_control_dependency(src_op, earliest_op, swap_in.op)
+            bw_frontier_ops -= fuse_bw_frontier_ops
+        return bw_frontier_ops
+
+    def _get_branch_ops(self, within_ops, threshold=0):
+        """Get ops whose order compared to the minimum order
+        is greater than the threshold.
+
+        Args:
+          within_ops: a set of `tf.Operation`.
+          threshold: an integer.
+
+        Return:
+          A set of `tf.Operation`.
+        """
+        orders = {self._topo_sort.get_order(op)
+                  for op in within_ops}
+        if not orders:
+            return set()
+        min_order = min(orders) + threshold
+        branch_ops = {
+            op
+            for op in within_ops
+            if (self._topo_sort.get_order(op) > min_order)}
+        return branch_ops
+
+    def _insert_swap_nodes(self, src_op):
+        """Insert swapin and swapout ops for the given operation into the graph.
+
+        This method does an in-place modification to the graph.
+
+        Args:
+          src_op: a `tf.Operation`
+        """
+        self._log_info("Operation: {}".format(src_op), 2)
+
+        # bypass excluded ops
+        if src_op in self._excl_ops:
+            return
+
+        # if inclusive mode is enabled, only proceed if this op is included
+        if self._incl_ops:
+            if src_op not in self._incl_ops:
+                return
+
+        for t in src_op.outputs:
+            if self._swapped_max_tensors():
+                return
+
+            frontier_ops = set(util.get_consuming_ops(t))
+            self._log_info("my frontier ops: {}".format(frontier_ops), 2)
+
+            bw_frontier_ops = frontier_ops & self._grad_ops
+            self._log_info("my bw frontier ops: {}".format(bw_frontier_ops), 2)
+
+            # swap branch ops if they are far enough (depending on threshold)
+            if self._swap_branches:
+                fw_branch_ops = self._get_branch_ops(
+                    frontier_ops - self._grad_ops,
+                    self._branch_threshold)
+                bw_frontier_ops = bw_frontier_ops | fw_branch_ops
+
+            # Do not swap tensors used by bw ops without outgoing ops.
+            # These bw ops can be removed by Tensorflow compiler
+            bw_frontier_ops = {op
+                               for op in bw_frontier_ops
+                               if set(self._get_forward_walk_ops(op, inclusive=False))}
+
+            if not bw_frontier_ops:
+                continue
+
+            self._log_info("Operation: {}, order {}, type {}".format(
+                src_op.name, self._topo_sort.get_order(src_op),
+                src_op.type), 1)
+
+            # create swap_out node only if there exists a real dest. operation
+            swapout_op = None
+            for op in bw_frontier_ops:
+                if self._topo_sort.get_order(op) >= 0:
+                    swapout_op = self._add_swapout(src_op, t)
+                    self._incpu_count = self._incpu_count + 1
+                    break
+
+            # create swap_in nodes
+            if self._fuse_swapins and swapout_op:
+                bw_frontier_ops = self._fuse_swapin_ops(
+                    src_op, swapout_op, bw_frontier_ops, t)
+            for dest_op in bw_frontier_ops:
+                if self._topo_sort.get_order(dest_op) < 0:
+                    if src_op in self._grad_ops:
+                        continue
+                    else:
+                        new_src_ops = self._find_new_src_op(dest_op)
+                        for op in new_src_ops:
+                            self._insert_swap_nodes(op)
+                else:
+                    # swap_in op
+                    swapin_op = self._add_swapin(swapout_op, dest_op, t)
+                    # control dependency -> swap_in
+                    self._add_control_dependency(src_op, dest_op, swapin_op)
+
+    def _add_swapout(self, src_op, ts0):
+        """Add a swapout operation to the graph to swap out the output tensor `ts0`
+        of the operation `src_op`.
+
+        This method does an in-place modification to the graph.
+
+        Example: the graph before and after this method invoked.
+        ```
+        Before
+          (src_op) -> |ts0| -> (dest_op)
+
+        After:
+          (src_op) -> |ts0| -> (swapout_op)
+          |ts0| -> (dest_op)
+        ```
+
+        Args:
+          src_op: a `tf.Operation` that produces the tensor `ts0`.
+          ts0: a output `tf.Tensor` of `src_op` being swapped out.
+
+        Return:
+          A `tf.Operation` newly added to the graph.
+        """
+        with tf.device(self._cpu_device):
+            swap_out = tf.identity(ts0, name="lms/swapout")
+
+        # Connect: src-node -> swap-out
+        src_svg = ge.sgv(src_op, graph=self._graph)
+        src_out_idx = src_svg.output_index(ts0)
+        self._connect_ops(src_op, swap_out.op, remap_outputs=True,
+                          idx=src_out_idx)
+        self._excl_ops.add(swap_out.op)
+        self._log_info("Tensor {} will be placed on {}".format(
+            ts0.name, self._cpu_device), 1)
+
+        return swap_out.op
+
+    def _add_swapin(self, swapout_op, dest_op, ts0):
+        """Add a swapin operation to the graph. The swapin ops reads
+        the output tensor of `swapout_op` and passes it to `dest_op`,
+        replacing the input tensor `ts0` of `dest_op`.
+
+        This method does an in-place modification to the graph.
+
+        Example: the graph before and after this method invoked.
+        ```
+        Before
+          |ts0| -> (swapout_op)
+          |ts0| -> (dest_op)
+
+        After:
+          |ts0| -> (swapout_op) -> (swapin_op) -> (dest_op)
+        ```
+
+        Args:
+          swapout_op: a `tf.Operation` that swapped out the tensor `ts0`.
+          dest_op: a `tf.Operation` that will consume the output tensor of `swapout_op`.
+          ts0: a `tf.Tensor` being the original input tensor of `dest_op`.
+
+        Return:
+          A `tf.Operation` newly added to the graph.
+        """
+        with tf.device(self._cpu_device):
+            swap_in = tf.identity(ts0, name="lms/swapin")
+
+        # Connect: swap_out -> swap_in
+        self._connect_ops(swapout_op, swap_in.op)
+
+        # Connect: swap_in -> dest
+        dest_svg = ge.sgv(dest_op, graph=self._graph)
+        input_idx = dest_svg.input_index(ts0)
+        self._connect_ops(swap_in.op, dest_op, remap_inputs=True, idx=input_idx)
+        self._excl_ops.add(swap_in.op)
+
+        self._log_info("Consuming op {} (order {}) swaps in {}".format(
+            dest_op.name, self._topo_sort.get_order(dest_op),
+            ts0.name), 1)
+
+        return swap_in.op
+    
+    def _add_swapin_str(self, swapout_op, dest_ops, ts0):
+        """Add a swapin operation to the graph. The swapin ops reads
+        the output tensor of `swapout_op` and passes it to `dest_op`,
+        replacing the input tensor `ts0` of `dest_op`.
+
+        Return:
+          A `tf.Operation` newly added to the graph.
+        """
+        with tf.device(self._cpu_device):
+            swap_in = tf.identity(ts0, name="lms-str/swapin")
+
+        # Connect: swap_out -> swap_in
+        self._connect_ops(swapout_op, swap_in.op)
+
+        for dest_op in dest_ops:
+            dest_svg = ge.sgv(dest_op, graph=self._graph)
+            if ts0 not in dest_svg._input_ts:
+                self._log_info("Tensor {} is not the input of svg {}, skip linking!".format(ts0, dest_svg))
+                return None
+            input_idx = dest_svg.input_index(ts0)
+            self._connect_ops(swap_in.op, dest_op, remap_inputs=True, idx=input_idx)
+
+        self._excl_ops.add(swap_in.op)
+        self._log_info("Consuming op {}, swaps in {}".format(
+            dest_ops, ts0.name), 1)
+
+        return swap_in.op
+
+    def _add_control_dependency(self, fw_op, bw_op, swapin_op):
+        """Find and add a control dependency to the graph.
+
+        This method does an in-place modification to the graph.
+
+        Args:
+          fw_op: a `tf.Operation`.
+          bw_op: a `tf.Operation`.
+          swapin_op: a `tf.Operation`.
+        """
+        # if lb is out of range, reset it to make sure
+        # that a control dependency op will be found
+        lb = self._lb
+        if (self._topo_sort.get_order(bw_op) - lb <=
+                self._topo_sort.get_order(fw_op)):
+            lb = 1
+        if fw_op in self._grad_ops:
+            re = self._do_direct_order(fw_op, bw_op, lb, self._ub)
+        elif self._ctrld_strategy is CTRLD_Strategy.CHAIN_RULE:
+            re = self._do_chain_rule(fw_op, bw_op, lb, self._ub)
+        elif self._ctrld_strategy is CTRLD_Strategy.DIRECT_ORDER:
+            re = self._do_direct_order(fw_op, bw_op, lb, self._ub)
+        else:
+            re = self._do_chain_rule(fw_op, bw_op, lb, self._ub)
+
+        ctrld_op = re[0]
+        ctrld_order = re[1]
+        if ctrld_op:
+            ge.add_control_inputs(swapin_op, ctrld_op)
+            self._log_info(
+                "Control dependency op {},  order: {}".format(
+                    ctrld_op.name, ctrld_order), 1)
+        else:
+            self._log_info(
+                "No control dependency op needed for swap in of op {}.".format(
+                    fw_op.name), 1)
+
+    def _find_new_src_op(self, original_op):
+        """Find a set of new operations to swap out their output tensors.
+
+        This method is used when `original_op` produces a tensor that is consumed by
+        a backward ops whose order is negative. In this case, the tensor might be consumed
+        immediately by the backward ops, depending on TensorFlow runtime. Hence, there is
+        no need to swap out the tensor.
+
+        This method starts from `original_op` and returns operations whose output tensors
+        are consumed by backward operations with positive order.
+
+        Args:
+          `original_op`: a `tf.Operation`.
+
+        Return:
+          A set of `tf.Operation`.
+        """
+        src_ops = set()
+        open_set = Queue.Queue()
+        closed_set = set()
+
+        open_set.put(original_op)
+
+        while not open_set.empty():
+            src_op = open_set.get()
+
+            # do action for src_op
+            next_ops = set()
+
+            frontier_ops = set()
+            for t in src_op.outputs:
+                frontier_ops |= set(util.get_consuming_ops(t))
+            has_order_ops = {
+                op
+                for op in frontier_ops
+                if (self._topo_sort.get_order(op) >
+                    self._topo_sort.bw_starting_order)
+            }
+            if has_order_ops:
+                src_ops.add(src_op)
+
+            next_ops = frontier_ops - has_order_ops
+            for op in next_ops:
+                if op in closed_set:
+                    continue
+                if op not in open_set.queue:
+                    open_set.put(op)
+
+            closed_set.add(src_op)
+        return src_ops
+
+    def _do_chain_rule(self, fw_op, bw_op, lower_b, upper_b):
+        """Find a control dependency operation using chain rules.
+        Go down along the forward phase to find corresponding backward ops
+        as candidates for control dependency ops.
+
+        Args:
+          fw_op: a `tf.Operation` that has a tensor swapped out.
+          bw_op: a `tf.Operation` that consumes a tensor swapped in.
+          lower_b: an `integer`. The distance in the graph between
+            `fw_op` and a forward operation that has corresponding backward
+            ops as candidates for control dependency ops must be greater than
+            `lower_b`.
+          upper_b: an `integer`. The distance in the graph between
+            `fw_op` and a forward operation that has corresponding backward
+             ops as candidates for control dependency ops must be smaller than
+            `upper_b`
+
+        Return:
+          A tuple of (`tf.Operation`, an `integer`). The first item is
+          the control dependency operation that triggers swapping in the input
+          tensor of `bw_op`. The second item is the order of the control
+          dependency operation in the topological order.
+        """
+        fw_order = self._topo_sort.get_order(fw_op)
+        bw_order = self._topo_sort.get_order(bw_op)
+
+        # check if the bw op is near the boundary between fw and bw phases
+        if (bw_order - lower_b) < self._topo_sort.bw_starting_order:
+            return self._do_direct_order(fw_op, bw_op, lower_b, upper_b)
+
+        open_set1 = Queue.Queue()
+        open_set2 = Queue.Queue()
+        closed_set = set()
+
+        open_set1.put(fw_op)
+
+        result_ops = set()
+        while not open_set1.empty():
+            # stop if reaching the upperbound
+            if upper_b == 0 or (lower_b > upper_b):
+                break
+
+            src_op = open_set1.get()
+
+            # do action for src_op
+            total_consumming_ops = set()
+            for t in src_op.outputs:
+                consumming_ops = set(util.get_consuming_ops(t))
+                total_consumming_ops |= consumming_ops
+
+            if lower_b <= 0:
+                # inside the range
+                consumming_ops_bw = total_consumming_ops & self._grad_ops
+                # check validation
+                consumming_ops_bw = {
+                    op
+                    for op in consumming_ops_bw
+                    if self._topo_sort.get_order(op) > fw_order}
+                consumming_ops_bw = {
+                    op
+                    for op in consumming_ops_bw
+                    if self._topo_sort.get_order(op) < bw_order}
+                consumming_ops_bw = {
+                    op
+                    for op in consumming_ops_bw
+                    if "/cond/" not in op.name}
+                result_ops |= consumming_ops_bw
+            # go to the next level
+            next_ops = total_consumming_ops - self._grad_ops
+            for op in next_ops:
+                if op in closed_set:
+                    continue
+                if op not in open_set2.queue:
+                    open_set2.put(op)
+
+            closed_set.add(src_op)
+            if open_set1.empty():
+                if result_ops:
+                    break
+                lower_b = lower_b - 1
+                upper_b = upper_b - 1
+                while not open_set2.empty():
+                    open_set1.put(open_set2.get())
+        if result_ops:
+            ctrld_op = next(iter(result_ops))
+            return (ctrld_op, self._topo_sort.get_order(ctrld_op))
+        else:
+            return (None, -1)
+
+    def _do_direct_order(self, fw_op, src_op, lower_b, upper_b):
+        """Find a control dependency operation using topological sort.
+
+        Args:
+          fw_op: a `tf.Operation` that has a tensor swapped out.
+          bw_op: a `tf.Operation` that consumes a tensor swapped in.
+          lower_b: an `integer`. The distance in the topological order
+            between `bw_op` and a candidate for control dependency ops
+            must be greater than `lower_b`.
+          upper_b: an `integer`. The distance in the topological order
+            between `bw_op` and a candidate for control dependency ops
+            must be smaller than `upper_b`
+
+        Return:
+          A tuple of (`tf.Operation`, an `integer`). The first item is
+          the control dependency operation that triggers swapping in the input
+          tensor of `bw_op`. The second item is the order of the control
+          dependency operation in the topological order.
+        """
+        result_ops = set()
+
+        # offset ordering
+        fw_order = self._topo_sort.get_order(fw_op)
+        src_order = self._topo_sort.get_order(src_op)
+
+        range_ub = src_order - lower_b
+        range_lb = max([src_order - upper_b, fw_order]) + 1
+
+        ctrld_order = -1
+        for i in reversed(range(range_lb, range_ub)):
+            candidates = self._topo_sort.get_ops(i)
+            # on the chain rule path
+            candidates = {op
+                          for op in candidates
+                          if src_op in set(self._get_forward_walk_ops(op))}
+            candidates = {op
+                          for op in candidates
+                          if "/cond/" not in op.name}
+            if candidates:
+                result_ops |= candidates
+                ctrld_order = i
+                break
+
+        if result_ops:
+            ctrld_op = next(iter(result_ops))
+            return (ctrld_op, ctrld_order)
+        else:
+            return (None, -1)
+
+    def _log_info(self, message, level=0):
+        """Log debug information.
+
+        Args:
+          message: a formatted string.
+          level: an `integer`.
+        """
+        if level == 0 or (self._debug and self._debug_level >= level):
+            # Use tf.logging.info instead of print, since print
+            # is not thread safe, which can break tests.
+            tf.logging.info("[LMS][{}] {}".format(level, message))
+
+    def _print_configuration(self):
+        """Print configuration information about LMS.
+        """
+        if self._n_tensors == 0:
+            self._log_info("n_tensors: all tensors")
+        else:
+            self._log_info("n_tensors: {}".format(self._n_tensors))
+        self._log_info("lb: {}".format(self._lb))
+
+    def _connect_ops(self, src_op, dest_op, remap_inputs=False,
+                     remap_outputs=False, idx=None, disconnect_first=False):
+        """A wrapper of `tensorflow.contrib.graph_editor.connect`.
+
+        This method does an in-place modification to the graph.
+
+        Args:
+          src_op: a `tf.Operation`.
+          dest_op: a `tf.Operation`.
+          remap_inputs: remap the input of `dest_op` or not.
+          remap_outputs: remap the output of `src_op` or not.
+          idx: index of input or output tensor.
+          disconnect_first: True means the current outputs of sgv0 are
+            disconnected.
+        """
+        src_sgv = ge.sgv(src_op, graph=self._graph)
+        dest_sgv = ge.sgv(dest_op, graph=self._graph)
+        if remap_outputs:
+            src_sgv = src_sgv.remap_outputs([idx])
+        if remap_inputs:
+            dest_sgv = dest_sgv.remap_inputs([idx])
+
+        ge.connect(src_sgv, dest_sgv, disconnect_first)
+
+    def _swapped_max_tensors(self):
+        """Check whether we swapped enough tensors or not.
+        """
+        return ((self._n_tensors > 0) and
+                (self._incpu_count >= self._n_tensors))
+
+
+class LMSSessionRunHook(tf.train.SessionRunHook):
+    ''' This hook is to modify the input graph for Large Model Support
+    by adding swap operations.
+    '''
+    def __init__(self, optimizer_scopes, **kwargs):
+        """Create an LMSHook object to edit the graph for supporting large model.
+
+        Args:
+          optimizer_scopes: a set of scopes for the optimizers/solvers.
+          kwargs: the kwargs to pass to LMS. Note, the `graph` argument is
+                  removed from the kwargs before initializing LMS because
+                  the graph is obtained automatically by the SessionRunHook and
+                  is generally not available at hook initilization time.
+        """
+        kwargs.pop('graph', None)
+        self.lms_obj = LMS(optimizer_scopes, **kwargs)
+
+    def begin(self):
+        self.lms_obj.run(tf.get_default_graph())
+
+
+class LMSKerasCallback(tf.keras.callbacks.Callback):
+    """This callback is to modify the input graph for Large Model Support
+    during Keras training / fit by adding swap operations.
+    """
+
+    def __init__(self, optimizer_scopes_override=None, **kwargs):
+        """Create an LMSKerasCallback object to edit the graph for
+           supporting large model tensor swapping when using TensorFlow Keras.
+
+        Args:
+          optimizer_scopes_override: by default the LMSKerasCallback will
+                automatically discover the optimizer scopes from the Keras
+                model. This parameter allows overriding that automatic discovery
+                with a set of optimizer scope names.
+          kwargs: the kwargs to pass to LMS. Note, the `graph` argument is
+                  removed from the kwargs and not used for initializing LMS
+                  because the graph is obtained automatically by the
+                  Keras callback during the set_model method.
+        """
+        self._optimizer_scopes = optimizer_scopes_override
+        self._lms_args = kwargs
+        self._lms_args.pop('graph', None)
+
+    def set_model(self, model):
+        self.model = model
+        optimizer_scopes = self._optimizer_scopes
+        if not self._optimizer_scopes:
+            optimizer_name = self.model.optimizer.__class__.__name__
+            # TensorFlow pre-1.14 uses the 'training/' name scope.
+            # TensorFlow 1.14 optimizer operations do not have the prepended
+            # 'training/'.
+            optimizer_scopes = {'training/'+optimizer_name+'/gradients',
+                                optimizer_name+'/gradients'}
+
+        # Check if the model has the train_function created. If the train
+        # function is created, (Keras fit, fit_generator, TensorFlow Keras fit)
+        # paths, then the optimizer operations / backward phase operations are
+        # in the graph.
+        if getattr(model, 'train_function') is None:
+            # This is the tf.keras fit_generator path for TensorFlow 1.14.
+            # The train_function has not been created, the graph is not
+            # "complete" yet because it will not have the optimizer and backward
+            # phases in it. We will create the train function now
+            # so the model is fully populated for running LMS on it.
+            model._make_train_function()
+
+        lmsMod = LMS(optimizer_scopes,
+                     graph=tf.get_default_graph(),
+                     **self._lms_args)
+        lmsMod.run()
diff --git a/tensorflow/python/ops/stropt/memory_saving_gradients.py b/tensorflow/python/ops/stropt/memory_saving_gradients.py
new file mode 100644
index 00000000..10ce1c1d
--- /dev/null
+++ b/tensorflow/python/ops/stropt/memory_saving_gradients.py
@@ -0,0 +1,422 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import time
+import sys
+import contextlib
+import numpy as np
+from toposort import toposort
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.keras import backend
+from tensorflow.contrib import graph_editor as ge
+from tensorflow.python.platform import tf_logging as logging
+sys.setrecursionlimit(10000)
+# refers back to current module if we decide to split helpers out
+util = sys.modules[__name__]
+
+# getting rid of "WARNING:tensorflow:VARIABLES collection name is deprecated"
+setattr(ops.GraphKeys, "VARIABLES", "variables")
+
+# save original gradients since tf.gradient could be monkey-patched to point
+# to our version
+from tensorflow.python.ops import gradients as tf_gradients_lib
+tf_gradient_function = tf_gradients_lib.gradients
+
+# ISSUE: https://github.com/cybertronai/gradient-checkpointing/issues/38
+def tf_gradients(ys, *args, **kwargs):
+    """Decorate tf.gradients calls with explicit device placement to avoid memory
+    leaks when splitting model across multiple GPUs"""
+    source = ys[0] if isinstance(ys, (list, tuple)) else ys
+    device = source.op.node_def.device if isinstance(source, ops.Tensor) else None
+    with ops.device(device):
+        return tf_gradient_function(ys, *args, **kwargs)
+
+MIN_CHECKPOINT_NODE_SIZE=1024    # use lower value during testing
+
+# specific versions we can use to do process-wide replacement of tf.gradients
+def gradients_speed(ys, xs, grad_ys=None, **kwargs):
+    return gradients(ys, xs, grad_ys, checkpoints='speed', **kwargs)
+
+def gradients_memory(ys, xs, grad_ys=None, **kwargs):
+    return gradients(ys, xs, grad_ys, checkpoints='memory', **kwargs)
+        
+def gradients_collection(ys, xs, grad_ys=None, **kwargs):
+    return gradients(ys, xs, grad_ys, checkpoints='collection', **kwargs)
+
+def gradients(ys, xs, grad_ys=None, checkpoints='collection', **kwargs):
+    '''
+    Authors: Tim Salimans & Yaroslav Bulatov
+
+    memory efficient gradient implementation inspired by "Training Deep Nets with Sublinear Memory Cost"
+    by Chen et al. 2016 (https://arxiv.org/abs/1604.06174)
+
+    ys,xs,grad_ys,kwargs are the arguments to standard tensorflow tf.gradients
+    (https://www.tensorflow.org/versions/r0.12/api_docs/python/train.html#gradients)
+
+    'checkpoints' can either be
+        - a list consisting of tensors from the forward pass of the neural net
+          that we should re-use when calculating the gradients in the backward pass
+          all other tensors that do not appear in this list will be re-computed
+        - a string specifying how this list should be determined. currently we support
+            - 'speed':  checkpoint all outputs of convolutions and matmuls. these ops are usually the most expensive,
+                        so checkpointing them maximizes the running speed
+                        (this is a good option if nonlinearities, concats, batchnorms, etc are taking up a lot of memory)
+            - 'memory': try to minimize the memory usage
+                        (currently using a very simple strategy that identifies a number of bottleneck tensors in the graph to checkpoint)
+            - 'collection': look for a tensorflow collection named 'checkpoints', which holds the tensors to checkpoint
+    '''
+
+    if not isinstance(ys,list):
+        ys = [ys]
+    if not isinstance(xs,list):
+        xs = [xs]
+    # traverse from backward and forward to remove meaningless nodes
+    bwd_ops = ge.get_backward_walk_ops([y.op for y in ys],
+                                       inclusive=True)
+
+    debug_print("bwd_ops: %s", bwd_ops)
+    
+    # forward ops are all ops that are candidates for recomputation
+    fwd_ops = ge.get_forward_walk_ops([x.op for x in xs],
+                                      inclusive=True,
+                                      within_ops=bwd_ops)
+    debug_print("fwd_ops: %s", fwd_ops)
+
+    # exclude ops with no inputs
+    fwd_ops = [op for op in fwd_ops if op.inputs]
+
+    # don't recompute xs, remove variables
+    xs_ops = _to_ops(xs)
+    fwd_ops = [op for op in fwd_ops if not op in xs_ops]
+    fwd_ops = [op for op in fwd_ops if not '/assign' in op.name]
+    fwd_ops = [op for op in fwd_ops if not '/Assign' in op.name]
+    fwd_ops = [op for op in fwd_ops if not '/read' in op.name]
+    ts_all = ge.filter_ts(fwd_ops, True) # get the tensors
+    ts_all = [t for t in ts_all if '/read' not in t.name]
+    ts_all = set(ts_all) - set(xs) - set(ys)
+
+    # construct list of tensors to checkpoint during forward pass, if not
+    # given as input
+    if type(checkpoints) is not list:
+        if checkpoints == 'collection':
+            checkpoints = ops.get_collection('checkpoints')
+            
+        elif checkpoints == 'speed':
+            # checkpoint all expensive ops to maximize running speed
+            checkpoints = ge.filter_ts_from_regex(fwd_ops, 'conv2d|Conv|MatMul')
+            
+        elif checkpoints == 'memory':
+            # remove very small tensors and some weird ops
+            def fixdims(t): # tf.Dimension values are not compatible with int, convert manually
+                try:
+                    return [int(e if e.value is not None else 64) for e in t]
+                except:
+                    return [0]  # unknown shape
+            ts_all = [t for t in ts_all if np.prod(fixdims(t.shape)) > MIN_CHECKPOINT_NODE_SIZE]
+            ts_all = [t for t in ts_all if 'L2Loss' not in t.name]
+            ts_all = [t for t in ts_all if 'entropy' not in t.name]
+            ts_all = [t for t in ts_all if 'FusedBatchNorm' not in t.name]
+            ts_all = [t for t in ts_all if 'Switch' not in t.name]
+            ts_all = [t for t in ts_all if 'dropout' not in t.name]
+            # DV: FP16_FIX - need to add 'Cast' layer here to make it work for FP16
+            ts_all = [t for t in ts_all if 'Cast' not in t.name]
+
+            # capture backward tensors in bwd_ops
+            with util.capture_ops() as bwd_ops:
+                tf_gradients(ys, xs, grad_ys, **kwargs)
+            debug_print("bwd_ops after call the gradient function: %s", bwd_ops)
+
+            bwd_inputs = [t for op in bwd_ops for t in op.inputs]
+            # list of tensors in forward graph that is in input to bwd graph
+            ts_filtered = list(set(bwd_inputs).intersection(ts_all))
+            debug_print("Using tensors %s", ts_filtered)
+            
+            # try two slightly different ways of getting bottlenecks tensors
+            # to checkpoint
+            for ts in [ts_filtered, ts_all]:
+
+                # get all bottlenecks in the graph
+                bottleneck_ts = []
+                for t in ts:
+                    b = set(ge.get_backward_walk_ops(t.op, inclusive=True, within_ops=fwd_ops))
+                    f = set(ge.get_forward_walk_ops(t.op, inclusive=False, within_ops=fwd_ops))
+                    # check that there are not shortcuts
+                    b_inp = set([inp for op in b for inp in op.inputs]).intersection(ts_all)
+                    f_inp = set([inp for op in f for inp in op.inputs]).intersection(ts_all)
+                    # if str(t.name).__contains__("dense/Relu"):
+                    if not set(b_inp).intersection(f_inp) and len(b_inp)+len(f_inp) >= len(ts_all):
+                        bottleneck_ts.append(t)  # we have a bottleneck!
+                    else:
+                        debug_print("Rejected bottleneck candidate and ops %s", [t] + list(set(ts_all) - set(b_inp) - set(f_inp)))
+
+                # success? or try again without filtering?
+                if len(bottleneck_ts) >= np.sqrt(len(ts_filtered)): # yes, enough bottlenecks found!
+                    break
+
+            if not bottleneck_ts:
+                raise Exception('unable to find bottleneck tensors! please provide checkpoint nodes manually, or use checkpoints="speed".')
+
+            # sort the bottlenecks
+            bottlenecks_sorted_lists = tf_toposort(bottleneck_ts, within_ops=fwd_ops)
+            sorted_bottlenecks = [t for ts in bottlenecks_sorted_lists for t in ts]
+
+            # save an approximately optimal number ~ sqrt(N)
+            N = len(ts_filtered)
+            if len(bottleneck_ts) <= np.ceil(np.sqrt(N)):
+                checkpoints = sorted_bottlenecks
+            else:
+                step = int(np.ceil(len(bottleneck_ts) / np.sqrt(N)))
+                checkpoints = sorted_bottlenecks[step::step]
+            
+        else:
+            raise Exception('%s is unsupported input for "checkpoints"' % (checkpoints,))
+
+    checkpoints = list(set(checkpoints).intersection(ts_all))
+
+    # at this point automatic selection happened and checkpoints is list of nodes
+    assert isinstance(checkpoints, list)
+
+    # better error handling of special cases
+    # xs are already handled as checkpoint nodes, so no need to include them
+    xs_intersect_checkpoints = set(xs).intersection(set(checkpoints))
+    if xs_intersect_checkpoints:
+        debug_print("Warning, some input nodes are also checkpoint nodes: %s",
+                    xs_intersect_checkpoints)
+    ys_intersect_checkpoints = set(ys).intersection(set(checkpoints))
+    logging.info("ys: %s, checkpoints: %s, intersect: %s", ys, checkpoints,
+                ys_intersect_checkpoints)
+    # saving an output node (ys) gives no benefit in memory while creating
+    # new edge cases, exclude them
+    if ys_intersect_checkpoints:
+        debug_print("Warning, some output nodes are also checkpoints nodes: %s",
+              format_ops(ys_intersect_checkpoints))
+
+    # remove initial and terminal nodes from checkpoints list if present
+    checkpoints = list(set(checkpoints) - set(ys) - set(xs))
+    # check that we have some nodes to checkpoint
+    if not checkpoints:
+        raise Exception('no checkpoints nodes found or given as input! ')
+
+    # mock, only leave one checkpoint
+    # checkpoints = [checkpoints[1], checkpoints[2]]
+    ###################################
+
+    debug_print(f"Number of checkpoints: {len(checkpoints)}:")
+    debug_print(f"{checkpoints}\n\n")
+
+    # 1. Disconnect dependencies between checkpointed tensors gradient func 
+    #    will not add gradient node for x, but the identical output.
+    #    `grad_node` can be used for further calc.
+    checkpoints_disconnected = {}
+    for x in checkpoints:
+        if x.op and x.op.name is not None:
+            grad_node = array_ops.stop_gradient(x, name=x.op.name+"_sg")
+        else:
+            grad_node = array_ops.stop_gradient(x)
+        grad_node.op._set_device(x.op.node_def.device)
+        checkpoints_disconnected[x] = grad_node
+    debug_print(f"checkpoints_disconnected: {checkpoints_disconnected}")
+    # visit back from loss node to checkpointed node, here builds sgv for the last checkpoint boundary
+    ops_to_copy = fast_backward_ops(seed_ops=[y.op for y in ys],
+                                    stop_at_ts=checkpoints, within_ops=fwd_ops)
+    debug_print("Found %s ops to copy within fwd_ops %s, seed %s, stop_at %s",
+                    len(ops_to_copy), fwd_ops, [r.op for r in ys], checkpoints)
+    debug_print("ops_to_copy = %s", ops_to_copy)
+    debug_print("Processing list %s", ys)
+
+    # 2. SubGraphView will build a subgraph with `ops_to_copy`.
+    #    copy_with_input_replacements(ops_to_copy, {}) breaks all input tensors for this subgraph, and
+    #    returns backward nodes under the gradient scope!
+    copied_sgv, info = ge.copy_with_input_replacements(ge.sgv(ops_to_copy), {})
+    for origin_op, op in info._transformed_ops.items():
+        op._set_device(origin_op.node_def.device) # grad op is set to co-located with its forward op
+    copied_ops = info._transformed_ops.values()
+    debug_print("Copied_svg: %s", copied_sgv)
+    debug_print("Copied %s to %s", ops_to_copy, copied_ops)
+
+    # 3. This rewiring makes checkpointed values point to subgraphs, which will will be treated 
+    #    as inputs of this subgraph. reroute_ts(ts0, ts1), break the link between ts1 and downstream op,
+    #    and insert ts0.
+    ge.reroute_ts(checkpoints_disconnected.values(), checkpoints_disconnected.keys(), can_modify=copied_ops)
+    debug_print("Rewired %s in place of %s restricted to %s",
+                checkpoints_disconnected.values(), checkpoints_disconnected.keys(), copied_ops)
+
+    # get gradients with respect to current boundary + original x's
+    copied_ys = [info._transformed_ops[y.op]._outputs[0] for y in ys]
+    boundary = list(checkpoints_disconnected.values()) # boundaries identify using which checkpoint to recompute
+    dv = tf_gradients(ys=copied_ys, xs=boundary+xs, grad_ys=grad_ys, **kwargs)
+    debug_print("Got gradients %s", dv)
+    debug_print("for %s", copied_ys)
+    debug_print(f"with respect to boundary:{boundary}, xs:{xs}")
+    
+    inputs_to_do_before = [y.op for y in ys]
+    if grad_ys is not None:
+        inputs_to_do_before += grad_ys
+    wait_to_do_ops = list(copied_ops) + [g.op for g in dv if g is not None]
+    my_add_control_inputs(wait_to_do_ops, inputs_to_do_before)
+
+    # partial derivatives to the checkpointed nodes
+    # dictionary of "node: backprop" for nodes in the boundary
+    d_checkpoints = {r: dr for r,dr in zip(checkpoints_disconnected.keys(),
+                                        dv[:len(checkpoints_disconnected)])}
+    # partial derivatives to xs (usually the params of the neural net)
+    d_xs = dv[len(checkpoints_disconnected):]
+    debug_print(f"Checkpoint gradients before loop {d_checkpoints}")
+    # 4. Loop the above procedure for each checkpoint, i.e., Algo.2 in paper
+    checkpoints_sorted_lists = tf_toposort(checkpoints, within_ops=fwd_ops)
+    debug_print(f"checkpoints_sorted_lists: {checkpoints_sorted_lists}")
+    for ts in checkpoints_sorted_lists[::-1]:
+        debug_print("Processing list %s", ts)
+        checkpoints_other = [r for r in checkpoints if r not in ts]
+        checkpoints_disconnected_other = [checkpoints_disconnected[r] for r in checkpoints_other]
+
+        # Reversely visiting forward nodes from current checkpoint, stopping at the former checkpoint
+        ops_to_copy = fast_backward_ops(within_ops=fwd_ops, seed_ops=[r.op for r in ts], stop_at_ts=checkpoints_other)
+        debug_print("Found %s ops to copy within %s, seed %s, stop_at %s",
+                    len(ops_to_copy), fwd_ops, [r.op for r in ts],
+                    checkpoints_other)
+        debug_print("ops_to_copy = %s", ops_to_copy)
+        if not ops_to_copy: # we're done!
+            break
+        copied_sgv, info = ge.copy_with_input_replacements(ge.sgv(ops_to_copy), {})
+        for origin_op, op in info._transformed_ops.items():
+            op._set_device(origin_op.node_def.device)
+        copied_ops = info._transformed_ops.values()
+        debug_print("Copied %s to %s", ops_to_copy, copied_ops)
+        ge.reroute_ts(checkpoints_disconnected_other, checkpoints_other, can_modify=copied_ops)
+        debug_print("Rewired %s in place of %s restricted to %s",
+                    checkpoints_disconnected_other, checkpoints_other, copied_ops)
+
+        # current checkpoint's output as new boundary
+        boundary = [info._transformed_ops[r.op]._outputs[0] for r in ts]
+        # the grad. of the last boundary as sources of backprop
+        substitute_backprops = [d_checkpoints[r] for r in ts]
+        dv = tf_gradients(boundary,
+                          checkpoints_disconnected_other+xs,
+                          grad_ys=substitute_backprops, **kwargs)
+        debug_print("Got gradients %s", dv)
+        debug_print("for %s", boundary)
+        debug_print("with respect to %s", checkpoints_disconnected_other+xs)
+        debug_print("with boundary backprop substitutions %s", substitute_backprops)
+
+        inputs_to_do_before = [d_checkpoints[r].op for r in ts]
+        wait_to_do_ops = list(copied_ops) + [g.op for g in dv if g is not None]
+        my_add_control_inputs(wait_to_do_ops, inputs_to_do_before)
+
+        # partial derivatives to the checkpointed nodes
+        for r, dr in zip(checkpoints_other, dv[:len(checkpoints_other)]):
+            if dr is not None:
+                if d_checkpoints[r] is None:
+                    d_checkpoints[r] = dr
+                else:
+                    d_checkpoints[r] += dr
+        def _unsparsify(x):
+            if not isinstance(x, ops.IndexedSlices):
+                return x
+            assert x.dense_shape is not None, "memory_saving_gradients encountered sparse gradients of unknown shape"
+            indices = x.indices
+            while indices.shape.ndims < x.values.shape.ndims:
+                indices = gen_array_ops.expand_dims(indices, -1)
+            return gen_array_ops.scatter_nd(indices, x.values, x.dense_shape)
+
+        # partial derivatives to xs (usually the params of the neural net)
+        d_xs_new = dv[len(checkpoints_other):]
+        for j in range(len(xs)):
+            if d_xs_new[j] is not None:
+                if d_xs[j] is None:
+                    d_xs[j] = _unsparsify(d_xs_new[j])
+                else:
+                    d_xs[j] += _unsparsify(d_xs_new[j])
+
+
+    return d_xs
+
+def tf_toposort(ts, within_ops=None):
+    all_ops = ge.get_forward_walk_ops([x.op for x in ts], within_ops=within_ops)
+
+    deps = {}
+    for op in all_ops:
+        for o in op.outputs:
+            deps[o] = set(op.inputs)
+    sorted_ts = toposort(deps)
+
+    # only keep the tensors from our original list
+    ts_sorted_lists = []
+    for l in sorted_ts:
+        keep = list(set(l).intersection(ts))
+        if keep:
+            ts_sorted_lists.append(keep)
+
+    return ts_sorted_lists
+
+def fast_backward_ops(within_ops, seed_ops, stop_at_ts):
+    bwd_ops = set(ge.get_backward_walk_ops(seed_ops, stop_at_ts=stop_at_ts))
+    ops = bwd_ops.intersection(within_ops).difference([t.op for t in stop_at_ts])
+    return list(ops)
+
+@contextlib.contextmanager
+def capture_ops():
+  """Decorator to capture ops created in the block.
+  with capture_ops() as ops:
+    # create some ops
+  print(ops) # => prints ops created.
+  """
+
+  micros = int(time.time()*10**6)
+  scope_name = str(micros)
+  op_list = []
+  with ops.name_scope(scope_name):
+    yield op_list
+
+  g = ops.get_default_graph()
+  op_list.extend(ge.select_ops(scope_name+"/.*", graph=g))
+
+def _to_op(tensor_or_op):
+  if hasattr(tensor_or_op, "op"):
+    return tensor_or_op.op
+  return tensor_or_op
+
+def _to_ops(iterable):
+  if not _is_iterable(iterable):
+    return iterable
+  return [_to_op(i) for i in iterable]
+
+def _is_iterable(o):
+  try:
+    _ = iter(o)
+  except Exception:
+    return False
+  return True
+
+DEBUG_LOGGING=True
+def debug_print(s, *args):
+  """Like logger.log, but also replaces all TensorFlow ops/tensors with their
+  names. Sensitive to value of DEBUG_LOGGING, see enable_debug/disable_debug
+
+  Usage:
+    debug_print("see tensors %s for %s", tensorlist, [1,2,3])
+  """
+
+  if DEBUG_LOGGING:
+    formatted_args = [format_ops(arg) for arg in args]
+    print("DEBUG "+s % tuple(formatted_args))
+
+def format_ops(ops, sort_outputs=True):
+  """Helper method for printing ops. Converts Tensor/Operation op to op.name,
+  rest to str(op)."""
+    
+  if hasattr(ops, '__iter__') and not isinstance(ops, str):
+    l = [(op.name if hasattr(op, "name") else str(op)) for op in ops]
+    if sort_outputs:
+      return sorted(l)
+    return l
+  else:
+    return ops.name if hasattr(ops, "name") else str(ops)
+
+def my_add_control_inputs(wait_to_do_ops, inputs_to_do_before):
+    for op in wait_to_do_ops:
+        ci = [i for i in inputs_to_do_before if op.control_inputs is None or i not in op.control_inputs]
+        ge.add_control_inputs(op, ci)
diff --git a/tensorflow/python/ops/stropt/str_strategy.py b/tensorflow/python/ops/stropt/str_strategy.py
new file mode 100644
index 00000000..3708cabd
--- /dev/null
+++ b/tensorflow/python/ops/stropt/str_strategy.py
@@ -0,0 +1,313 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import time
+from toposort import toposort
+from typing import List
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops.stropt import memory_saving_gradients
+from tensorflow.python.ops.stropt.strategy_parser import SolvedStrategy
+from tensorflow.python.ops.stropt.lms import LMS
+from tensorflow.contrib import graph_editor as ge
+from tensorflow.python.ops import gradients as tf_gradients_lib
+tf_gradient_function = tf_gradients_lib.gradients
+
+
+########
+VERBOSE = True
+########
+
+def verbose_log(log):
+    if VERBOSE:
+        print("[STR DEBUG] " + str(log))
+
+def naive_checkpointing(ys,
+                        xs,
+                        grad_ys=None):
+    """ Customized gradient function.
+    """
+    with ops.get_default_graph()._mutation_lock():
+        verbose_log("Pure recomputing strategy with Chen's greedy strategy")
+        return memory_saving_gradients.gradients_memory(ys, xs, grad_ys)
+
+def str_gradients(ys,
+                xs,
+                grad_ys=None,
+                config=None, optimizer_name="", **kwargs):
+    """ Edit the graph for recomputation during the gradient calculation.
+    """
+    
+    if not isinstance(ys,list):
+        ys = [ys]
+    if not isinstance(xs,list):
+        xs = [xs]
+
+    verbose_log("Parsing STR configuration: {}".format([dict(config.items(section)) for section in config.sections()]))
+    solution_type = config["solution"]["strategy"]
+
+    if solution_type not in ["recompute", "swap", "hybrid"]:
+        verbose_log("Solution type {} doesn't has corresponding optimizations, run normal training".format(solution_type))
+        return tf_gradient_function(ys, xs, grad_ys)
+
+    global VERBOSE
+    VERBOSE = True if config["solution"]["verbose"] == "true" else False
+    
+    results = config[solution_type] # r,p,q
+    params = dict(results)
+    params["layer_names"] = config["solution"]["tags"]
+    ss = SolvedStrategy(params)
+
+    with ops.get_default_graph()._mutation_lock():
+        if solution_type == "recompute":
+            if hasattr(ss, "r"):
+                verbose_log("Use recomputing with Checkmate's solution")
+                return edit_recompute(ss=ss, ys=ys, xs=xs, grad_ys=None, **kwargs)
+            else:
+                verbose_log("Solution is not configured, use Chen's heuristic to select checkpoints")
+                return naive_checkpointing(ys=ys, xs=xs, grad_ys=None)
+        elif solution_type == "swap":
+            verbose_log("Use swap strategy of DYNPROG")
+            # Calc. gradients normally
+            grads = tf_gradient_function(ys, xs, grad_ys)
+            edit_swap(ss, optimizer_name)
+            return grads
+        elif solution_type == "hybrid":
+            if len(ss.r.recompute.keys()) == 0:
+                verbose_log("No ops need to be recomputed!")
+                # Calc. gradients normally
+                grads = tf_gradient_function(ys, xs, grad_ys)
+            else:
+                # Add recompute nodes and gradient nodes in the default graph
+                grads = edit_recompute(ss=ss, ys=ys, xs=xs, grad_ys=None, **kwargs)
+                verbose_log("Finish editing for recomputation!")
+            # Add swap nodes according to node names
+            edit_swap(ss, optimizer_name)
+            verbose_log("Finish editing for swapping!")
+            return grads
+        else:
+            pass
+
+def edit_recompute(ss: SolvedStrategy, 
+                ys=None,
+                xs=None,
+                grad_ys=None, **kwargs):
+    """ This part is modified from cybertronai's gradient-checkpointing.
+    """
+    verbose_log(f"Edit graph for recomputation ops")
+    if not isinstance(ys,list):
+        ys = [ys]
+    if not isinstance(xs,list):
+        xs = [xs]
+    # traverse from backward and forward to filter out minor ops
+    bwd_ops = ge.get_backward_walk_ops([y.op for y in ys],
+                                       inclusive=True)
+    fwd_ops = ge.get_forward_walk_ops([x.op for x in xs],
+                                      inclusive=True,
+                                      within_ops=bwd_ops)
+    # exclude ops with no inputs
+    fwd_ops = [op for op in fwd_ops if op.inputs]
+    xs_ops = _to_ops(xs)
+    fwd_ops = [op for op in fwd_ops if not op in xs_ops]
+    fwd_ops = [op for op in fwd_ops if not '/assign' in op.name]
+    fwd_ops = [op for op in fwd_ops if not '/Assign' in op.name]
+    fwd_ops = [op for op in fwd_ops if not '/read' in op.name]
+    fwd_ops = [op for op in fwd_ops if not '/Read' in op.name]
+    ts_all = ge.filter_ts(fwd_ops, True) # get the tensors
+    ts_all = [t for t in ts_all if '/read' not in t.name]         
+    # remove parameters and the loss
+    ts_all = set(ts_all) - set(xs) - set(ys)
+    
+    sorted_fwd = sorted_fwd_ops(fwd_ops)
+    verbose_log(f"Sorted forward ops: {fwd_ops}")
+    # Identify tensors that need to be recomputed
+    recomp_ops_list = ss.get_ops_by_layer(list(ss.r.recompute.keys()), sorted_fwd)
+
+    verbose_log(f"Ops to be recomputed: {recomp_ops_list}")
+    checkpoints = ss.get_checkpoint_tensors(recomp_ops_list, sorted_fwd)
+    verbose_log(f"Supplementary tensors to be checkpoints: {checkpoints}")
+    # 1. Disconnect dependencies between checkpointed tensors gradient func 
+    #    will not add gradient node for x, but the identical output.
+    #    `grad_node` can be used for further calc.
+    checkpoints_disconnected = {}
+    for x in checkpoints:
+            if x.op and x.op.name is not None:
+                grad_node = array_ops.stop_gradient(x, name=x.op.name+"_sg")
+            else:
+                grad_node = array_ops.stop_gradient(x)
+            grad_node.op._set_device(x.op.node_def.device)
+            checkpoints_disconnected[x] = grad_node
+    verbose_log(f"Checkpointed tensors and distconnected: {checkpoints_disconnected}")
+
+    # visit back from loss node to checkpointed node, here builds sgv for the last checkpoint boundary
+    ops_to_copy = fast_backward_ops(seed_ops=[y.op for y in ys],
+                                    stop_at_ts=checkpoints, within_ops=fwd_ops)
+    verbose_log(f"Found {len(ops_to_copy)} ops to copy within fwd_ops {fwd_ops}, seed {[r.op for r in ys]}, stop_at {checkpoints}")
+    verbose_log(f"ops_to_copy = {ops_to_copy}")
+    verbose_log(f"Processing list {ys}")
+    
+    copied_sgv, info = ge.copy_with_input_replacements(ge.sgv(ops_to_copy), {})
+    for origin_op, op in info._transformed_ops.items():
+        op._set_device(origin_op.node_def.device) # grad op is set to co-located with its forward op
+    copied_ops = info._transformed_ops.values()
+    verbose_log(f"Copied_svg: {copied_sgv}")
+    verbose_log(f"Copied {ops_to_copy} to {copied_ops}")
+
+    ge.reroute_ts(checkpoints_disconnected.values(), checkpoints_disconnected.keys(), can_modify=copied_ops)
+    verbose_log(f"Rewired {checkpoints_disconnected.values()} in place of {checkpoints_disconnected.keys()} restricted to {copied_ops}")
+    
+    # get gradients with respect to current boundary + original x's
+    # copied_ys = copied_svg.outputs[0]
+    copied_ys = [info._transformed_ops[y.op]._outputs[0] for y in ys]
+    boundary = list(checkpoints_disconnected.values()) # boundaries identify using which checkpoint to recompute
+    dv = tf_gradient_function(ys=copied_ys, xs=boundary+xs, grad_ys=grad_ys, **kwargs)
+    verbose_log(f"Got gradients {dv}")
+    verbose_log(f"for {copied_ys}")
+    verbose_log(f"with respect to ops, boundary:{boundary}, xs:{xs}")
+    
+    inputs_to_do_before = [y.op for y in ys]
+    if grad_ys is not None:
+        inputs_to_do_before += grad_ys
+    wait_to_do_ops = list(copied_ops) + [g.op for g in dv if g is not None]
+    batch_add_control_inputs(wait_to_do_ops, inputs_to_do_before)
+
+    # partial derivatives to the checkpointed nodes
+    # dictionary of "node: backprop" for nodes in the boundary
+    d_checkpoints = {r: dr for r,dr in zip(checkpoints_disconnected.keys(),
+                                        dv[:len(checkpoints_disconnected)])}
+    # partial derivatives to xs (usually the params+baise of the neural net)
+    d_xs = dv[len(checkpoints_disconnected):]
+    verbose_log(f"dv of checkpoints: {d_checkpoints}")
+    
+    checkpoints_sorted_lists = tf_toposort(checkpoints, within_ops=fwd_ops)
+    verbose_log(f"checkpoints_sorted_lists: {checkpoints_sorted_lists}")
+    # 4. Loop the above procedure for each checkpoint, i.e., Algo.2 in paper
+    for ts in checkpoints_sorted_lists[::-1]:
+        verbose_log("=========loop========")
+        verbose_log(f"Processing list {ts}")
+        checkpoints_other = [r for r in checkpoints if r not in ts]
+        checkpoints_disconnected_other = [checkpoints_disconnected[r] for r in checkpoints_other]
+
+        # Reversely visiting forward nodes from current checkpoint, stopping at the former checkpoint
+        ops_to_copy = fast_backward_ops(within_ops=fwd_ops, seed_ops=[t.op for t in ts], stop_at_ts=checkpoints_other)
+        verbose_log(f"Found {len(ops_to_copy)} ops to copy within {fwd_ops}, seed {[t.op for t in ts]}, stop_at {checkpoints_other}")
+        verbose_log(f"ops_to_copy = {ops_to_copy}")
+
+        if not ops_to_copy: # we're done!
+            break
+
+        copied_sgv, info = ge.copy_with_input_replacements(ge.sgv(ops_to_copy), {})
+        for origin_op, op in info._transformed_ops.items():
+            op._set_device(origin_op.node_def.device)
+        copied_ops = info._transformed_ops.values()
+        verbose_log(f"Copied {ops_to_copy} to {copied_ops}")
+        ge.reroute_ts(checkpoints_disconnected_other, checkpoints_other, can_modify=copied_ops)
+        verbose_log(f"Rewired {checkpoints_disconnected_other} in place of {checkpoints_other} restricted to {copied_ops}")
+
+        # current checkpoint's output as new boundary
+        boundary = [info._transformed_ops[r.op]._outputs[0] for r in ts]
+        # the grad. of the last boundary as sources of backprop
+        substitute_backprops = [d_checkpoints[r] for r in ts]
+        dv = tf_gradient_function(boundary,
+                          checkpoints_disconnected_other+xs,
+                          grad_ys=substitute_backprops, **kwargs)
+        verbose_log(f"Got gradients {dv}")
+        verbose_log(f"for {boundary}")
+        verbose_log(f"with respect to boundary: {checkpoints_disconnected_other}, xs: {xs}")
+        verbose_log(f"with boundary backprop substitutions {substitute_backprops}")
+
+
+        inputs_to_do_before = [d_checkpoints[r].op for r in ts]
+        wait_to_do_ops = list(copied_ops) + [g.op for g in dv if g is not None]
+        batch_add_control_inputs(wait_to_do_ops, inputs_to_do_before)
+
+        # partial derivatives to the checkpointed nodes
+        for r, dr in zip(checkpoints_other, dv[:len(checkpoints_other)]):
+            if dr is not None:
+                if d_checkpoints[r] is None:
+                    d_checkpoints[r] = dr
+                else:
+                    d_checkpoints[r] += dr
+        def _unsparsify(x):
+            if not isinstance(x, ops.IndexedSlices):
+                return x
+            assert x.dense_shape is not None, "memory_saving_gradients encountered sparse gradients of unknown shape"
+            indices = x.indices
+            while indices.shape.ndims < x.values.shape.ndims:
+                indices = gen_array_ops.expand_dims(indices, -1)
+            return gen_array_ops.scatter_nd(indices, x.values, x.dense_shape)
+
+        # partial derivatives to xs (usually the params of the neural net)
+        d_xs_new = dv[len(checkpoints_other):]
+        for j in range(len(xs)):
+            if d_xs_new[j] is not None:
+                if d_xs[j] is None:
+                    d_xs[j] = _unsparsify(d_xs_new[j])
+                else:
+                    d_xs[j] += _unsparsify(d_xs_new[j])
+        verbose_log("=========end========")
+    return d_xs
+
+def sorted_fwd_ops(base_fwd_ops):
+    """ sort base_fwd_ops according to the order in the graph order"""
+    sorted_fwd_ops =ge.make_list_of_op(ops.get_default_graph())
+    ret = [op for op in sorted_fwd_ops if op in base_fwd_ops]
+    return ret
+
+def fast_backward_ops(within_ops, seed_ops, stop_at_ts):
+    bwd_ops = set(ge.get_backward_walk_ops(seed_ops, stop_at_ts=stop_at_ts))
+    ops = bwd_ops.intersection(within_ops).difference([t.op for t in stop_at_ts])
+    return list(ops)
+
+def batch_add_control_inputs(wait_to_do_ops, inputs_to_do_before):
+    for op in wait_to_do_ops:
+        ci = [i for i in inputs_to_do_before if op.control_inputs is None or i not in op.control_inputs]
+        ge.add_control_inputs(op, ci)
+
+def tf_toposort(ts, within_ops=None):
+    all_ops = ge.get_forward_walk_ops([x.op for x in ts], within_ops=within_ops)
+
+    deps = {}
+    for op in all_ops:
+        for o in op.outputs:
+            deps[o] = set(op.inputs)
+    sorted_ts = toposort(deps)
+
+    # only keep the tensors from our original list
+    ts_sorted_lists = []
+    for l in sorted_ts:
+        keep = list(set(l).intersection(ts))
+        if keep:
+            ts_sorted_lists.append(keep)
+
+    return ts_sorted_lists
+
+def _to_op(tensor_or_op):
+  if hasattr(tensor_or_op, "op"):
+    return tensor_or_op.op
+  return tensor_or_op
+
+def _to_ops(iterable):
+  if not _is_iterable(iterable):
+    return iterable
+  return [_to_op(i) for i in iterable]
+
+def _is_iterable(o):
+  try:
+    _ = iter(o)
+  except Exception:
+    return False
+  return True
+
+# Implementation of edit swapping graph
+def edit_swap(ss: SolvedStrategy, optimizer_name: str):
+    optimizer_scopes = {'training/'+optimizer_name+'/gradients',
+                                optimizer_name+'/gradients'}
+    lmsMod = LMS(optimizer_scopes,
+                     graph=ops.get_default_graph(),
+                     solved_strategy=ss)
+    lmsMod.run()
+    
diff --git a/tensorflow/python/ops/stropt/strategy_parser.py b/tensorflow/python/ops/stropt/strategy_parser.py
new file mode 100644
index 00000000..62ce80d0
--- /dev/null
+++ b/tensorflow/python/ops/stropt/strategy_parser.py
@@ -0,0 +1,342 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from typing import Dict, List, Tuple
+import re
+from tensorflow.python.framework.ops import Operation, Tensor
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.contrib import graph_editor as ge
+import numpy as np
+
+
+VERBOSE = True
+def verbose_log(log):
+    if VERBOSE:
+        print("[STR DEBUG] " + str(log))
+
+supported_form = ["p", "q", "r", "layer_names"]
+layer_output_tag = ["/Relu", "/Softmax", "/Tanh", "/Sigmoid", "/Softplus", "/Softsign", "/Selu", "/Elu"]
+backward_ops_tag = "training/"
+
+class Matrix:
+    """ This class handles the matrix parsing and format transformation.
+    """
+    def __init__(self, path: str, name):
+        assert name in supported_form, "Unknown matrices"
+        
+        self.m_name = name
+        self.m = np.loadtxt(path, delimiter=" ")
+
+        # {layer_id: List[recompute_stage]}, recompute_stages
+        self.recompute = None
+        # {layer_id: List[swap_out_start_stage]}
+        self.swap_out = None
+        # {layer_id: List[swap_in_start_stage]}
+        self.swap_in = None
+
+    def __str__(self) -> str:
+        if self.m_name == "r":
+            return "Recompute: {}".format(str(self.recompute))
+        elif self.m_name == "p":
+            return "Swap-out: {}".format(str(self.swap_out))
+        else:
+            return "Swap-in: {}".format(str(self.swap_in))
+
+class SwapRecord:
+    def __init__(self, layer_id, layer_tag=None) -> None:
+        self.layer_id = layer_id
+        self.layer_tag = layer_tag
+
+        # A tensor to be swapped out
+        self.swap_out_source_ts = None
+        # Swap out operation will be a control input of swap_out_control_op
+        self.swap_out_control_op = None
+        # Swap in operations will be a control input of swap_in_control_ops
+        self.swap_in_control_ops = list()
+        # Ops which depends on the swapped tensor
+        self.swap_in_dest_ops = list()
+
+    def set_swap_out(self, ts, control_op):
+        if self.swap_out_source_ts is not None:
+            logging.warn("A swaprecord only support one swap-out source.")
+            return
+        self.swap_out_source_ts = ts
+        self.swap_out_control_op = control_op
+
+    def add_swap_in_op(self, in_control_op: Operation, dest_ops: List[Operation]):
+        # A swap-in tensor might has multiple destination operations
+        self.swap_in_control_ops.append(in_control_op)
+        self.swap_in_dest_ops.append(dest_ops)
+    
+    def __str__(self) -> str:
+        attrs = []
+        attrs.append("layer_id={}, layer_tag={}".format(self.layer_id, self.layer_tag))
+        attrs.append("out_source_ts={}".format([self.swap_out_source_ts]))
+        attrs.append("out_control_op={}".format([self.swap_out_control_op]))
+        attrs.append("in_control_ops={}".format(self.swap_in_control_ops))
+        attrs.append("in_dest_ops={}".format(self.swap_in_dest_ops))
+        return ", ".join(attrs)
+
+
+class SolvedStrategy:
+    """ Form the optimization strategy.
+    """  
+    def __init__(self, strategy: Dict):
+        if "layer_names" in strategy.keys():
+            # The layer name needs contain the operation scope name, which can be used
+            # to select operations within this layer.
+            self.layer_names = self.parse_layer("layer_names", strategy["layer_names"])
+
+            # Parsing matrices
+            if "r" in strategy.keys():
+                self.r = self.parsed_rpq("r", strategy["r"])
+            if "p" in strategy.keys():
+                self.p = self.parsed_rpq("p", strategy["p"])
+            if "q" in strategy.keys():
+                self.q = self.parsed_rpq("q", strategy["q"])
+        else:
+            logging.warn("Missing layer name file.")
+
+    def parsed_rpq(self, name, path) -> Matrix:
+        """ Construct matrix object.
+
+        Args:
+            name (str): matrix name
+            path (str): file path of the matrix
+
+        Returns:
+            Matrix: [description]
+        """        
+        mat = Matrix(path, name)
+        assert mat.m.shape[0] == (len(self.layer_names) * 2 - 1)
+        filled = dict()
+        # visit low-triangle
+        for j in range(0, mat.m.shape[0] - 1):
+            for i in range(j + 1, mat.m.shape[0]):
+                if mat.m[i, j] == 1:
+                    if j not in filled.keys():
+                        filled[j] = []
+                    filled[j].append(i)
+        if name == "r":
+            mat.recompute = filled
+        elif name == "p":
+            mat.swap_out = filled
+        else:
+            mat.swap_in = filled
+        return mat
+
+    
+    def parse_layer(self, name, path) -> List:
+        """ 
+            Ensure: len(layer_names)*2 - 1 == matrix.shape[0] == matrix.shape[1]
+
+        Returns:
+            List: [description]
+        """        
+        layer_names = list()
+        with open(path, "r") as flayer:
+            lines = flayer.readlines()
+            for line in lines:
+                layer_names.append(line.strip())
+        return layer_names
+
+
+    def get_ops_by_layer(self, layer_id, within_ops: List[Operation], contain_fwd=True, contain_bwd=True) -> List[Operation]:
+        """ Query operations in `within_ops` to find all operations within the `layer_id`.
+            The tag string of the scope of this layer is queried from self.layer_names[layer_id].
+
+        Args:
+            layer_id (str|List): id of layer/layers listed in matrix
+            within_ops (List): a sorted list of tf operators
+
+        Returns:
+            List[Operation]: Operations that need to be queried.
+        """                
+        condition = None
+        if contain_fwd and contain_bwd:
+            condition = 1 # all ops are allowed
+        elif contain_fwd and not contain_bwd:
+            condition = 2 # only forward
+        elif not contain_fwd and contain_bwd:
+            condition = 3 # only backward
+        else:
+            pass
+
+        def get_ops(layer_id):
+            if layer_id >= len(self.layer_names):
+                layer_id = self.bwd_to_fwd(layer_id)
+            layer_ops = list()
+            tag = self.layer_names[layer_id]
+            # walk_into = False
+            for op in within_ops:
+                if condition == 1 and \
+                    op.name.__contains__(tag):
+                    layer_ops.append(op)
+                elif condition == 2 and \
+                    op.name.__contains__(tag) and not op.name.__contains__(backward_ops_tag):
+                    layer_ops.append(op)
+                elif condition == 3 and \
+                    not op.name.__contains__(tag) and op.name.__contains__(backward_ops_tag):
+                    layer_ops.append(op)
+                else:
+                    pass
+                    # walk_into = True
+                # if walk_into and not str(op.name).__contains__(tag):
+                    # break
+            return layer_ops
+        
+        if not isinstance(layer_id, list):
+            layer_id = [layer_id]
+        r_ops = list()
+        for l in layer_id:
+            r_ops.extend(get_ops(l))
+        return r_ops
+        
+    def bwd_to_fwd(self, layer_id):
+        return (len(self.layer_names) - 1) * 2 - layer_id
+
+
+    def get_checkpoint_tensors(self, recomp_ops: List[Operation], all_ops: List[Operation]) -> List[Tensor]:
+        """ The layer that contains recomp_ops will not in the checkpoint list.
+            For a checkpoint layer, we only choose the output tensor of activation node as the `checkpoint tensor`.
+
+        Args:
+            recomp_ops (List[Operation]): ops for recomputation
+            all_ops (List[Operation]): topology sorted op list
+
+        Returns:
+            List[Tensor]: [description]
+        """       
+        # tags for excluding some nodes
+        exclude_tag = ["loss/"]
+        ctags = "|".join(layer_output_tag)
+        etags = "|".join(exclude_tag)
+
+        recomp_op_names = [op.name for op in recomp_ops]
+        checkpoints = []
+        for op in all_ops:
+            if op.name in recomp_op_names:
+                continue
+            if re.search(ctags, op.name) and not re.search(etags, op.name) and len(op.outputs) > 0:
+                checkpoints.append(op.outputs[0])
+        return checkpoints
+
+    def __get_activation_nodes(self, within_ops: List[Operation], all=False):
+        """ Get activation nodes among within_ops.
+        """        
+        if len(within_ops) == 1:
+            # Some layer may only have one op, such as pooling layer
+            return within_ops[0]
+        
+        outop_tag = "|".join(layer_output_tag)
+        ops = list()        
+        for op in within_ops:
+            if re.search(outop_tag, op.name):
+                if not all:
+                    return op
+                else:
+                    ops.append(op)
+        if len(ops) == 0:
+            return None
+        else:
+            return ops
+
+    def get_swap_record(self, fwd_ops: List[Operation], grad_ops: List[Operation]) -> List[SwapRecord]:
+        """ For each swapped tensor, parsing the swapping source and destination operations, 
+            and saving them as SwapRecord objects.
+
+        Args:
+            fwd_ops (List[Operation]): topology sorted op list
+            grad_ops (List[Operation]): topology sorted op list
+
+        Returns:
+            List[SwapRecord]: [description]
+        """        
+        
+        rcds = list()
+        for layer_id in self.p.swap_out.keys():
+            if layer_id not in self.q.swap_in.keys():
+                verbose_log("Cannot find swap-out operations for swap-in at layer {}, skipping".format(layer_id))
+                continue
+            verbose_log("Processing layer {}'s swapping, swap out at {}, swap in at {}".format(layer_id, \
+                                                            self.p.swap_out[layer_id], self.q.swap_in[layer_id]))
+            # Query foward or backward ops to get the swapped op and its output tensors
+            if layer_id >= len(self.layer_names):
+                swaped_ops = self.get_ops_by_layer(layer_id, grad_ops)
+            else:
+                swaped_ops = self.get_ops_by_layer(layer_id, fwd_ops)
+            if len(swaped_ops) == 0:
+                verbose_log("Cannot find corresponding ops for layer_id {}".format(layer_id))
+                continue
+            # The output op of the swapped layer generates the swapped tensor.
+            swapped_op = self.__get_activation_nodes(swaped_ops)
+            if swapped_op is None:
+                verbose_log("Cannot find swap-out activation nodes for layer {}, among {}".format(layer_id, swaped_ops))
+                for op in swaped_ops[::-1]:
+                    if len(op.outputs) > 0:
+                        swapped_op = op
+                        break
+                verbose_log("Choose op that have outputs: {}".format(swapped_op))
+                # continue
+            swap_ts = swapped_op.outputs[0]
+
+            # Get swap out start op for controling, only one swap-out is allowed
+            out_stage_id = self.p.swap_out[layer_id][0]
+            if  out_stage_id >= len(self.layer_names):
+                swap_out_control_ops = self.get_ops_by_layer(out_stage_id, grad_ops)
+            else:
+                swap_out_control_ops = self.get_ops_by_layer(out_stage_id, fwd_ops)
+            # Choose the activation op, the swap_out op will be a control input of this op
+            control_op = self.__get_activation_nodes(swap_out_control_ops)
+            if control_op is None:
+                verbose_log("Cannot find swap-out control nodes for layer {}, use the first op {}".format(layer_id, swap_out_control_ops[0]))
+                control_op = swap_out_control_ops[0]
+                # continue
+            verbose_log("Find swapout ops: {}, \n\tchoose swap tensor {}, \n\tfinish at the end of ops: {}" \
+                                                                    .format(swapped_op, swap_ts, [control_op]))
+
+            record = SwapRecord(layer_id)
+            # Swap ts at the first op of swap_out_stage
+            record.set_swap_out(ts=swap_ts, control_op=control_op)
+
+            # Swap in may be triggered multiple times
+            in_stage_ids = self.q.swap_in[layer_id]
+            for in_id in in_stage_ids:
+                if in_id >= len(self.layer_names):
+                    swap_in_control_ops = self.get_ops_by_layer(in_id, grad_ops)
+                else:
+                    swap_in_control_ops = self.get_ops_by_layer(in_id, fwd_ops)
+                if not len(swap_in_control_ops) > 0:
+                    continue
+                # Choose the last op to control the swap finishing
+                swap_in_control_op = self.__get_activation_nodes(swap_in_control_ops)
+                # verbose_log("Swap in will finish at op: {}".format([swap_in_control_op]))
+                if swap_in_control_op is None:
+                    verbose_log("Cannot find swap-in control nodes for ref. layer {}, use the last op {}".format(in_id, swap_in_control_ops[-1]))
+                    swap_in_control_op = swap_in_control_ops[-1]
+                    # continue
+                consuming_ops = ge._util.get_consuming_ops(swap_ts)
+                # TODO: maybe have branch nodes during forward
+                consuming_ops = list(set(consuming_ops).intersection(set(grad_ops)))
+                # verbose_log("Get grad. consumer of this swap out in grad_ops: {}".format(consuming_ops))
+            
+                # Only the swap-in which has a downstream consumer is valid.
+                ops_fwd_from_in_control = ge.get_forward_walk_ops(seed_ops=swap_in_control_op, \
+                                                                inclusive=False, within_ops=fwd_ops + grad_ops)
+                valid_consume_op = []
+                for op in consuming_ops:
+                    # if op in ops_fwd_from_in_control or op.name.__contains__("_sg"):
+                    if op in ops_fwd_from_in_control:
+                        # for recompute
+                        valid_consume_op.append(op)
+                
+                if len(valid_consume_op) == 0:
+                    continue
+                record.add_swap_in_op(swap_in_control_op, valid_consume_op)
+                verbose_log("Add swap-in op at {} for consumer {}".format(swap_in_control_op, valid_consume_op))
+            if len(record.swap_in_dest_ops) == 0:
+                # Skip record that there are no destinations
+                continue
+            rcds.append(record)
+        return rcds
diff --git a/tensorflow/python/ops/stropt/topos.py b/tensorflow/python/ops/stropt/topos.py
new file mode 100644
index 00000000..3d792827
--- /dev/null
+++ b/tensorflow/python/ops/stropt/topos.py
@@ -0,0 +1,186 @@
+# (C) Copyright IBM Corp. 2018. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""TOPOS
+"""
+from six.moves import queue as Queue
+import toposort
+
+import tensorflow.contrib.graph_editor as ge
+from tensorflow.contrib.graph_editor import util
+
+
+class TOPOS(object):
+    """TOPOS class builds a topological order from the computational graph.
+    """
+    def __init__(self, seed_ops, grad_ops):
+        """Create a TOPOS object.
+
+        Args:
+          seed_ops: a list of `tf.Operation`.
+          grad_ops: a set of `tf.Operation`.
+        """
+        self._seed_ops = seed_ops
+        self._grad_ops = grad_ops
+
+        # index: set(ops)
+        self._topo_sort = {}
+        # op: index
+        self._orders = {}
+        self._bw_starting_order = -1
+
+    def build(self):
+        """Build a topological order
+        """
+        topo_sort = list(toposort.toposort(self._build_dependency_dict()))
+        for i in range(0, len(topo_sort)):
+            self._topo_sort[i] = topo_sort[i]
+
+        # if a bw op has the same order with a fw op,
+        # then remove the bw op
+        self._clean_bw_ops()
+
+        # if there are non-bw ops in the bw phase,
+        # then remove them, and do reordering
+        self._clean_update_ops()
+        self._reindex()
+
+        # build a dict of (op, order)
+        self._build_order_dict()
+
+        # starting order of the backward phase
+        for i in range(0, len(self._topo_sort)):
+            ops = self._topo_sort[i]
+            if (ops & self._grad_ops):
+                self._bw_starting_order = i
+                break
+
+    def _build_dependency_dict(self):
+        """Build a dictionary of dependencies among nodes.
+        """
+        open_set = Queue.Queue()
+        closed_set = set()
+
+        dep_dict = {}
+        for op in self._seed_ops:
+            open_set.put(op)
+
+        reachable_ops = set(ge.get_walks_intersection_ops(
+            list(self._seed_ops), list(self._grad_ops)))
+
+        # traversal in the fw phase
+        while not open_set.empty():
+            src_op = open_set.get()
+
+            # do action for src_op
+            dep_ops = set(src_op.control_inputs)
+            for t in src_op.inputs:
+                dep_ops |= set(util.get_generating_ops(t))
+                dep_ops &= reachable_ops
+            dep_dict[src_op] = dep_ops
+
+            next_ops = set()
+            for t in src_op.outputs:
+                next_ops |= set(util.get_consuming_ops(t))
+            for op in next_ops:
+                if op in closed_set:
+                    continue
+                if op not in open_set.queue:
+                    open_set.put(op)
+
+            closed_set.add(src_op)
+
+        return dep_dict
+
+    def _build_order_dict(self):
+        """Build a dictionary to quickly find an order of an ops.
+        """
+        for order, dep_ops in self._topo_sort.items():
+            for op in dep_ops:
+                self._orders[op] = order
+
+    def _clean_bw_ops(self):
+        """There are some bw ops that
+             - have no incoming bw ops except its fw op, or
+             - have no outgoing ops.
+        Execution order of these ops may depend on Tensorflow runtime.
+        """
+        for i in range(0, len(self._topo_sort)):
+            dep_ops = self._topo_sort[i]
+            fw_dep_ops = dep_ops - self._grad_ops
+            if fw_dep_ops:
+                self._topo_sort[i] = fw_dep_ops
+            else:
+                self._topo_sort[i] = dep_ops
+
+    def _clean_update_ops(self):
+        """Remove ops that are in the update phase.
+        """
+        update_ops = set(ge.get_forward_walk_ops(
+            list(self._grad_ops), inclusive=False))
+        for i in range(0, len(self._topo_sort)):
+            ops = self._topo_sort[i]
+            # remove ops that are not bw or fw op
+            # e.g ops in the update phase
+            self._topo_sort[i] = ops - update_ops
+
+    def _reindex(self):
+        """Remove orders with empty set and _reindex.
+        """
+        topo_sort = {}
+        index = 0
+        for i in range(0, len(self._topo_sort)):
+            ops = self._topo_sort[i]
+            if ops:
+                topo_sort[index] = ops
+                index += 1
+        self._topo_sort = topo_sort
+
+    def get_order(self, op):
+        """Return the order of an operation.
+
+        Args:
+          op: a `tf.Operation`.
+
+        Return:
+          An integer.
+        """
+        if op in self._orders:
+            return self._orders[op]
+        else:
+            return -1
+
+    def get_ops(self, order):
+        """Return a set of ops with the same order.
+
+        Args:
+          order: an integer.
+
+        Return:
+          A set of `tf.Operation`
+        """
+        return self._topo_sort[order]
+
+    @property
+    def size(self):
+        """The number of orders in the topological order.
+        """
+        return len(self._topo_sort)
+
+    @property
+    def bw_starting_order(self):
+        """The starting order of the backward phase.
+        """
+        return self._bw_starting_order