Rename wrapper to emitter, format code

Author: lantiga

torch2c/__init__.py

@@ -1,10 +1,12 @@
 import torch
 from torch.autograd import Variable
 import os
-from . import wrappers
+from . import emitters
+
 
 def _wrap(obj, prevfns=[]):
-    return wrappers._class_map[obj.__class__](obj,prevfns)
+    return emitters._class_map[obj.__class__](obj,prevfns)
+
 
 def _traverse_graph_recursive(out, el):
     if isinstance(el, Variable):
@@ -19,6 +21,7 @@ def _traverse_graph_recursive(out, el):
         for u in el.previous_functions:
             _traverse_graph_recursive(out,u[0])
 
+
 def _traverse_graph(node):
     nodes = []
     _traverse_graph_recursive(nodes,node.creator)
@@ -28,20 +31,22 @@ def _traverse_graph(node):
         el.infer_type(var_dict)
     return nodes
 
+
 def _wrap_out_node(nodes, out):
     out_node = _wrap(out)
     out_creator_id = id(out.creator)
     out_creator_node = [el for el in nodes if el.id == out_creator_id][0]
     out_node.infer_type({out_creator_id: out_creator_node})
     return out_node
 
+
 def _generate_c(nodes, out, fnname, out_path):
     # TODO: generate three functions
     # 1. load parameters (gen name from fnname)
     # 2. run forward
     # 3. free parameters (gen name from fnname)
     # to avoid loading from disk at each forward
-    var_nodes = [el for el in nodes if type(el) == wrappers.Variable]
+    var_nodes = [el for el in nodes if type(el) == emitters.Variable]
     out_node = _wrap_out_node(nodes,out)
     # TODO: make it more general re: last_node?
     last_node = nodes[-1]
@@ -61,20 +66,23 @@ def _generate_c(nodes, out, fnname, out_path):
     lines = [ifndef, includes, fndecl, '{'] + lines + ['}', endif]
     return '\n'.join(lines)
 
+
 def _to_persisted(var_node):
-    persisted = wrappers.PersistedVariable(var_node.obj,[])
+    persisted = emitters.PersistedVariable(var_node.obj,[])
     persisted.numtype = var_node.numtype
     return persisted
 
+
 def _clone_var(var):
     out = Variable(data=var.data.clone(),
                    creator=var.creator,
                    requires_grad=var.requires_grad,
                    volatile=var.volatile)
     return out
 
+
 def _generate_test(nodes, out, fnname, filename, out_path):
-    var_nodes = [_to_persisted(el) for el in nodes if type(el) == wrappers.Variable]
+    var_nodes = [_to_persisted(el) for el in nodes if type(el) == emitters.Variable]
     out_node = _to_persisted(_wrap_out_node(nodes,out))
     out_baseline_node = _to_persisted(_wrap_out_node(nodes,_clone_var(out)))
     out_node.obj.data.zero_()
@@ -93,6 +101,7 @@ def _generate_test(nodes, out, fnname, filename, out_path):
     lines = [includes, fndecl, '{'] + lines + ['}']
     return '\n'.join(lines)
 
+
 def compile(node, fnname, out_path, compile_test=False):
     nodes = _traverse_graph(node)
     if not os.path.isdir(out_path):
@@ -110,4 +119,3 @@ def compile(node, fnname, out_path, compile_test=False):
         with open(os.path.join(out_path,test_filename),'w') as f:
             f.write(test_src)
 
-

torch2c/wrappers.py renamed to torch2c/emitters.py

@@ -14,11 +14,12 @@
 
 _class_map = {}
 
-def register(wrapperClass, torchClass):
-    _class_map[torchClass] = wrapperClass
 
+def register(emitterClass, torchClass):
+    _class_map[torchClass] = emitterClass
 
-class Wrapper(object):
+
+class Emitter(object):
 
     def __init__(self, obj, prevfns):
         self.id = id(obj)
@@ -148,20 +149,21 @@ def tensor_meta_tpl(size_name, stride_name, size, stride=None):
 
 
 #####################
-# Wrapper subclasses
+# Emitter subclasses
 #####################
 
 
-class Variable(Wrapper):
+class Variable(Emitter):
 
     def __init__(self, obj, prevfns):
-        Wrapper.__init__(self, obj, prevfns)
+        Emitter.__init__(self, obj, prevfns)
 
     def infer_type(self, var_dict):
         self.numtype = self.obj.data.__class__.__name__[:len('Tensor')-1]
 
 register(Variable, torch.autograd.Variable)
 
+
 def persist_tensor(tensor, name, out_path, datadir, size_name='size_$id', stride_name='stride_$id'):
     contiguous = tensor.contiguous()
     filename = '%s.th' % name
@@ -177,6 +179,7 @@ def persist_tensor(tensor, name, out_path, datadir, size_name='size_$id', stride
     meta, meta_free = tensor_meta_tpl(size_name,stride_name,size,stride)
     return os.path.join(datadir,filename), meta, meta_free
 
+
 # TODO: add this function to an auxiliary file
 # call it something like TH${T}Storage_newFromFile(filename);
 def read_storage(storage_name,filepath,numtype):
@@ -205,6 +208,7 @@ def read_storage(storage_name,filepath,numtype):
     '''
     return Template(tpl).substitute(subs)
 
+
 class PersistedVariable(Variable):
 
     def __init__(self, obj, prevfns):
@@ -227,6 +231,7 @@ def free_tpl(self):
             TH${T}Storage_free(storage_$id);
             '''
 
+
 class Parameter(PersistedVariable):
 
     def __init__(self, obj, prevfns):
@@ -235,10 +240,10 @@ def __init__(self, obj, prevfns):
 register(Parameter, torch.nn.parameter.Parameter)
 
 
-class Linear(Wrapper):
+class Linear(Emitter):
 
     def __init__(self, obj, prevfns):
-        Wrapper.__init__(self, obj, prevfns)
+        Emitter.__init__(self, obj, prevfns)
 
         try:
             input, weight, bias = [id(el) for el in prevfns]
@@ -262,14 +267,13 @@ def free_tpl(self):
             TH${T}Tensor_free(addBuffer_$id);
             '''
 
-
 register(Linear, torch.nn._functions.linear.Linear)
 
 
-class LogSoftmax(Wrapper):
+class LogSoftmax(Emitter):
 
     def __init__(self, obj, prevfns):
-        Wrapper.__init__(self, obj, prevfns)
+        Emitter.__init__(self, obj, prevfns)
         self.def_vars({'input': id(prevfns[0])})
         self.infer_type_var = 'input'
 
@@ -287,10 +291,10 @@ def free_tpl(self):
 register(LogSoftmax, torch.nn._functions.thnn.auto.LogSoftmax)
 
 
-class Threshold(Wrapper):
+class Threshold(Emitter):
 
     def __init__(self, obj, prevfns):
-        Wrapper.__init__(self, obj, prevfns)
+        Emitter.__init__(self, obj, prevfns)
         self.def_vars({
             'input': id(prevfns[0]),
         })
@@ -315,10 +319,10 @@ def free_tpl(self):
 register(Threshold, torch.nn._functions.thnn.auto.Threshold)
 
 
-class Noop(Wrapper):
+class Noop(Emitter):
 
     def __init__(self, obj, prevfns):
-        Wrapper.__init__(self, obj, prevfns)
+        Emitter.__init__(self, obj, prevfns)
         self.def_vars({'input': id(prevfns[0])})
         self.infer_type_var = 'input'
 
@@ -334,10 +338,10 @@ def free_tpl(self):
 register(Noop, torch.nn._functions.dropout.FeatureDropout)
 
 
-class View(Wrapper):
+class View(Emitter):
 
     def __init__(self, obj, prevfns):
-        Wrapper.__init__(self, obj, prevfns)
+        Emitter.__init__(self, obj, prevfns)
         self.def_vars({'input': id(prevfns[0])})
         self.infer_type_var = 'input'
 
@@ -359,10 +363,10 @@ def free_tpl(self):
 register(View, torch.autograd._functions.tensor.View)
 
 
-class MaxPool2d(Wrapper):
+class MaxPool2d(Emitter):
 
     def __init__(self, obj, prevfns):
-        Wrapper.__init__(self, obj, prevfns)
+        Emitter.__init__(self, obj, prevfns)
         self.def_vars({
             'input': id(prevfns[0])
         })
@@ -392,10 +396,10 @@ def free_tpl(self):
 register(MaxPool2d, torch.nn._functions.thnn.pooling.MaxPool2d)
 
 
-class ConvNd(Wrapper):
+class ConvNd(Emitter):
 
     def __init__(self, obj, prevfns):
-        Wrapper.__init__(self, obj, prevfns)
+        Emitter.__init__(self, obj, prevfns)
         self.def_vars({'input': id(prevfns[0]),
                        'weight': id(prevfns[1]),
                        'bias': id(prevfns[2])})