user api to trigger partitioning

apache · Aug 14, 2019 · b6b5147 · b6b5147
1 parent 843c3ab
commit b6b5147
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/mxnet/symbol/symbol.py b/python/mxnet/symbol/symbol.py
@@ -1437,6 +1437,12 @@ def _gen_atomic_symbol(self):
         return Symbol(handle)
 
 
+    def optimizeFor(self, backend):
+        """Partition symbol and optimize it for a given backend"""
+        out = SymbolHandle()
+        check_call(_LIB.MXOptimizeForBackend(self.handle, c_str(backend), ctypes.byref(out)))
+
+
     # pylint: disable=too-many-locals
     def simple_bind(self, ctx, grad_req='write', type_dict=None, stype_dict=None,
                     group2ctx=None, shared_arg_names=None, shared_exec=None,