checked cast does it all

Author: killeent

Local.cwrap

@@ -80,11 +80,5 @@
     - TensorList tensors
     - int dim
   aten_custom_call: |
-    size_t inputs_ = tensors_.size();
-    std::vector<${THTensor}*> t_;
-    t_.reserve(inputs_);
-    for (unsigned int i = 0; i < inputs_; ++i) {
-      t_.push_back(tensors_[i]->tensor);
-    }
-    ${THTensor}_catArray(${state,}self_->tensor, t_.data(), inputs_, dim);
+    ${THTensor}_catArray(${state,}self_->tensor, tensors_.data(), tensors_.size(), dim);
 ]]

Utils.h

@@ -21,9 +21,10 @@ static inline T* checked_cast(Base* expr, const char * name, int pos) {
     T::typeString(),expr->type().toString(),pos,name);
 }
 
-template <typename T, typename TBase>
-static inline ArrayRef<T*> tensor_list_checked_cast(ArrayRef<TBase> tensors, const char * name, int pos) {
-  std::vector<T*> casted(tensors.size());
+// Converts a TensorList (i.e. ArrayRef<Tensor> to the underlying TH* Tensor Pointer)
+template <typename T, typename TBase, typename TH>
+static inline std::vector<TH*> tensor_list_checked_cast(ArrayRef<TBase> tensors, const char * name, int pos) {
+  std::vector<TH*> casted(tensors.size());
   for (unsigned int i = 0; i < tensors.size(); ++i) {
     auto *expr = tensors[i].pImpl;
     if (!expr) {
@@ -33,7 +34,7 @@ static inline ArrayRef<T*> tensor_list_checked_cast(ArrayRef<TBase> tensors, con
     }
     auto result = dynamic_cast<T*>(expr);
     if (result) {
-      casted.push_back(result);
+      casted.push_back(result->tensor);
     } else {
       runtime_error("Expected a Tensor of type %s but found a type %s for sequence element %u "
                     " in sequence argument at position #%d '%s'",

function_wrapper.py

@@ -99,7 +99,7 @@ def __init__(self, reason):
     'THStride*': CodeTemplate('THLongStorageView::make(${arg_name},true)'),
     'real': CodeTemplate('${arg_name}.to${ScalarName}()'),
     'accreal': CodeTemplate('${arg_name}.to${AccScalarName}()'),
-    'TensorList': CodeTemplate('tensor_list_checked_cast<${Tensor}>(${arg_name},"${arg_name}",${arg_pos})'),
+    'TensorList': CodeTemplate('tensor_list_checked_cast<${Tensor}, Tensor, ${THTensor}>(${arg_name},"${arg_name}",${arg_pos})'),
 }
 
 CHECKED_USE = {