Merge pull request pytorch#25 from killeent/nullable-tensors

add support for Null Tensors to functions

Author: killeent

Utils.h

@@ -10,8 +10,11 @@ namespace at {
 void runtime_error(const char *format, ...);
 
 template <typename T, typename Base>
-static inline T* checked_cast(Base* expr, const char * name, int pos) {
+static inline T* checked_cast(Base* expr, const char * name, int pos, bool allowNull) {
   if(!expr) {
+    if (allowNull) {
+      return (T*) expr;
+    }
     runtime_error("Expected a Tensor of type %s but found an undefined Tensor for argument #%d '%s'",
       T::typeString(),pos,name);
   }

common_with_cwrap.py

@@ -93,6 +93,14 @@ def signature(option, kwarg_only_count):
 
 def enumerate_options_due_to_default(declaration,
                                      allow_kwarg=True, type_to_signature=[], remove_self=True):
+
+    # Checks to see if an argument with a default keyword is a Tensor that
+    # by default can be NULL. In this case, instead of generating another
+    # option that excludes this argument, we will instead generate a single
+    # function call that allows for the Tensor to be NULL
+    def is_nullable_tensor_arg(arg):
+        return arg['type'] == 'THTensor*' and arg['default'] == 'nullptr'
+
     # TODO(zach): in cwrap this is shared among all declarations
     # but seems to assume that all declarations will have the same
     new_options = []

copy_wrapper.py

@@ -25,7 +25,7 @@
 FUNCTION = CodeTemplate("""\
 void ${Type}::copy(const Tensor & src, Tensor & dst) {
   // code generated by function_wrapper
-  auto dst_ = checked_cast<${Tensor}>(dst.pImpl,"dst",0);
+  auto dst_ = checked_cast<${Tensor}>(dst.pImpl,"dst",0,false);
   (void) dst_; //silence unused warning
   switch(src.type().ID()) {
     ${copy_body}

function_wrapper.py

@@ -109,11 +109,14 @@ def __init__(self, reason):
     'long': 'int64_t',
 }
 CHECKED_CAST = {
-    'THTensor*': CodeTemplate('checked_cast<${Tensor}>(${arg_name}.pImpl,"${arg_name}",${arg_pos})'),
-    'THBoolTensor*': CodeTemplate('checked_cast<${Backend}ByteTensor>(${arg_name}.pImpl,"${arg_name}",${arg_pos})'),
-    'THIndexTensor*': CodeTemplate('checked_cast<${Backend}LongTensor>(${arg_name}.pImpl,"${arg_name}",${arg_pos})'),
-    'THIntegerTensor*': CodeTemplate('checked_cast<${Backend}IntTensor>(${arg_name}.pImpl,"${arg_name}",${arg_pos})'),
-    'THStorage*': CodeTemplate('checked_cast<${Storage}>(&${arg_name},"${arg_name}",${arg_pos})'),
+    'THTensor*': CodeTemplate('checked_cast<${Tensor}>(${arg_name}.pImpl,"${arg_name}",${arg_pos}, ${null_okay})'),
+    'THBoolTensor*':
+        CodeTemplate('checked_cast<${Backend}ByteTensor>(${arg_name}.pImpl,"${arg_name}",${arg_pos}, ${null_okay})'),
+    'THIndexTensor*':
+        CodeTemplate('checked_cast<${Backend}LongTensor>(${arg_name}.pImpl,"${arg_name}",${arg_pos}, ${null_okay})'),
+    'THIntegerTensor*':
+        CodeTemplate('checked_cast<${Backend}IntTensor>(${arg_name}.pImpl,"${arg_name}",${arg_pos}, ${null_okay})'),
+    'THStorage*': CodeTemplate('checked_cast<${Storage}>(&${arg_name},"${arg_name}",${arg_pos}, false)'),
     'THGenerator*': CodeTemplate('check_generator(&${arg_name})'),
     'THSize*': CodeTemplate('THLongStorageView::make(${arg_name},true)'),
     'THStride*': CodeTemplate('THLongStorageView::make(${arg_name},true)'),
@@ -133,6 +136,8 @@ def __init__(self, reason):
     'TensorList': "{0}_.data(), {0}_.size()",
 }
 
+CHECKED_USE_NULLABLE = CodeTemplate('${arg_name}_ ? ${usage} : NULL')
+
 ALLOC_WRAP = {
     'THTensor*': 'new ${Tensor}(context)',
     'THBoolTensor*': 'new ${Backend}ByteTensor(context)',
@@ -341,12 +346,20 @@ def create_derived(backend_type_env, declarations):
     def requires_checked_cast(argument):
         return argument['type'] in CHECKED_CAST
 
+    def nullable_argument(argument):
+        return (argument['type'] == 'THTensor*' and
+                argument.get('default', '') == 'nullptr')
+
     def bool_option_is_string(argument):
         return 'if_true' in argument and isinstance(argument['if_true'], string_type)
 
     def get_argument(argument, option):
         if requires_checked_cast(argument):
-            return CHECKED_USE.get(argument['type'], '{}_').format(argument['name'])
+            checked_use = CHECKED_USE.get(argument['type'], '{}_').format(argument['name'])
+            if nullable_argument(argument):
+                checked_use = CHECKED_USE_NULLABLE.substitute(
+                                env={}, arg_name=argument['name'], usage=checked_use)
+            return checked_use
         elif argument['type'] == 'bool' and 'if_true' in argument:
             if bool_option_is_string(argument):
                 tpl = '({}) ? "{}" : "{}"'
@@ -424,8 +437,14 @@ def emit_body(env, option):
                         arg['name'], arg['name']))
                 # extract the TensorImpl from an existing tensor (or Storage, etc.)
                 else:
+                    # special case where we allow undefined Tensors, and thus
+                    # the checked cast succeeds even if the Tensor is not
+                    # defined
+                    null_okay = 'true' if nullable_argument(arg) else 'false'
+
                     check_cast = CHECKED_CAST[arg['type']].substitute(
-                        env, arg_name=arg['name'], arg_pos=count)
+                        env, arg_name=arg['name'], arg_pos=count,
+                        null_okay=null_okay)
                     body.append("auto {}_ = {};".format(
                         arg['name'], check_cast))
                 if drop_argument(arg, option):