torch-rb 0.6.0 → 0.8.2

data/ext/torch/templates.h

@@ -4,8 +4,7 @@
 #undef isfinite
 #endif
 
-#include <rice/Array.hpp>
-#include <rice/Object.hpp>
+#include <rice/rice.hpp>
 
 using namespace Rice;
 
@@ -22,6 +21,10 @@ using torch::IntArrayRef;
 using torch::ArrayRef;
 using torch::TensorList;
 using torch::Storage;
+using ScalarList = ArrayRef<Scalar>;
+
+using torch::nn::init::FanModeType;
+using torch::nn::init::NonlinearityType;
 
 #define HANDLE_TH_ERRORS                                             \
   try {
@@ -38,37 +41,42 @@ using torch::Storage;
 #define RETURN_NIL                                                   \
   return Qnil;
 
-template<>
-inline
-std::vector<int64_t> from_ruby<std::vector<int64_t>>(Object x)
-{
-  Array a = Array(x);
-  std::vector<int64_t> vec(a.size());
-  for (long i = 0; i < a.size(); i++) {
-    vec[i] = from_ruby<int64_t>(a[i]);
-  }
-  return vec;
-}
+class OptionalTensor {
+  torch::Tensor value;
+  public:
+    OptionalTensor(Object o) {
+      if (o.is_nil()) {
+        value = {};
+      } else {
+        value = Rice::detail::From_Ruby<torch::Tensor>().convert(o.value());
+      }
+    }
+    OptionalTensor(torch::Tensor o) {
+      value = o;
+    }
+    operator torch::Tensor() const {
+      return value;
+    }
+};
 
-template<>
-inline
-std::vector<Tensor> from_ruby<std::vector<Tensor>>(Object x)
+namespace Rice::detail
 {
-  Array a = Array(x);
-  std::vector<Tensor> vec(a.size());
-  for (long i = 0; i < a.size(); i++) {
-    vec[i] = from_ruby<Tensor>(a[i]);
-  }
-  return vec;
-}
+  template<>
+  struct Type<FanModeType>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
 
-class FanModeType {
-  std::string s;
+  template<>
+  class From_Ruby<FanModeType>
+  {
   public:
-    FanModeType(Object o) {
-      s = String(o).str();
-    }
-    operator torch::nn::init::FanModeType() {
+    FanModeType convert(VALUE x)
+    {
+      auto s = String(x).str();
       if (s == "fan_in") {
         return torch::kFanIn;
       } else if (s == "fan_out") {
@@ -77,22 +85,24 @@ class FanModeType {
         throw std::runtime_error("Unsupported nonlinearity type: " + s);
       }
     }
-};
-
-template<>
-inline
-FanModeType from_ruby<FanModeType>(Object x)
-{
-  return FanModeType(x);
-}
+  };
+
+  template<>
+  struct Type<NonlinearityType>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
 
-class NonlinearityType {
-  std::string s;
+  template<>
+  class From_Ruby<NonlinearityType>
+  {
   public:
-    NonlinearityType(Object o) {
-      s = String(o).str();
-    }
-    operator torch::nn::init::NonlinearityType() {
+    NonlinearityType convert(VALUE x)
+    {
+      auto s = String(x).str();
       if (s == "linear") {
         return torch::kLinear;
       } else if (s == "conv1d") {
@@ -119,102 +129,70 @@ class NonlinearityType {
         throw std::runtime_error("Unsupported nonlinearity type: " + s);
       }
     }
-};
+  };
+
+  template<>
+  struct Type<OptionalTensor>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
 
-template<>
-inline
-NonlinearityType from_ruby<NonlinearityType>(Object x)
-{
-  return NonlinearityType(x);
-}
+  template<>
+  class From_Ruby<OptionalTensor>
+  {
+  public:
+    OptionalTensor convert(VALUE x)
+    {
+      return OptionalTensor(x);
+    }
+  };
+
+  template<>
+  struct Type<Scalar>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
 
-class OptionalTensor {
-  torch::Tensor value;
+  template<>
+  class From_Ruby<Scalar>
+  {
   public:
-    OptionalTensor(Object o) {
-      if (o.is_nil()) {
-        value = {};
+    Scalar convert(VALUE x)
+    {
+      if (FIXNUM_P(x)) {
+        return torch::Scalar(From_Ruby<int64_t>().convert(x));
       } else {
-        value = from_ruby<torch::Tensor>(o);
+        return torch::Scalar(From_Ruby<double>().convert(x));
       }
     }
-    OptionalTensor(torch::Tensor o) {
-      value = o;
-    }
-    operator torch::Tensor() const {
-      return value;
+  };
+
+  template<typename T>
+  struct Type<torch::optional<T>>
+  {
+    static bool verify()
+    {
+      return true;
     }
-};
-
-template<>
-inline
-Scalar from_ruby<Scalar>(Object x)
-{
-  if (x.rb_type() == T_FIXNUM) {
-    return torch::Scalar(from_ruby<int64_t>(x));
-  } else {
-    return torch::Scalar(from_ruby<double>(x));
-  }
-}
-
-template<>
-inline
-OptionalTensor from_ruby<OptionalTensor>(Object x)
-{
-  return OptionalTensor(x);
-}
-
-template<>
-inline
-torch::optional<torch::ScalarType> from_ruby<torch::optional<torch::ScalarType>>(Object x)
-{
-  if (x.is_nil()) {
-    return torch::nullopt;
-  } else {
-    return torch::optional<torch::ScalarType>{from_ruby<torch::ScalarType>(x)};
-  }
-}
-
-template<>
-inline
-torch::optional<int64_t> from_ruby<torch::optional<int64_t>>(Object x)
-{
-  if (x.is_nil()) {
-    return torch::nullopt;
-  } else {
-    return torch::optional<int64_t>{from_ruby<int64_t>(x)};
-  }
-}
+  };
 
-template<>
-inline
-torch::optional<double> from_ruby<torch::optional<double>>(Object x)
-{
-  if (x.is_nil()) {
-    return torch::nullopt;
-  } else {
-    return torch::optional<double>{from_ruby<double>(x)};
-  }
-}
-
-template<>
-inline
-torch::optional<bool> from_ruby<torch::optional<bool>>(Object x)
-{
-  if (x.is_nil()) {
-    return torch::nullopt;
-  } else {
-    return torch::optional<bool>{from_ruby<bool>(x)};
-  }
-}
-
-template<>
-inline
-torch::optional<Scalar> from_ruby<torch::optional<Scalar>>(Object x)
-{
-  if (x.is_nil()) {
-    return torch::nullopt;
-  } else {
-    return torch::optional<Scalar>{from_ruby<Scalar>(x)};
-  }
+  template<typename T>
+  class From_Ruby<torch::optional<T>>
+  {
+  public:
+    torch::optional<T> convert(VALUE x)
+    {
+      if (NIL_P(x)) {
+        return torch::nullopt;
+      } else {
+        return torch::optional<T>{From_Ruby<T>().convert(x)};
+      }
+    }
+  };
 }

data/ext/torch/tensor.cpp

@@ -1,7 +1,6 @@
 #include <torch/torch.h>
 
-#include <rice/Constructor.hpp>
-#include <rice/Module.hpp>
+#include <rice/rice.hpp>
 
 #include "tensor_functions.h"
 #include "ruby_arg_parser.h"
@@ -11,6 +10,39 @@
 using namespace Rice;
 using torch::indexing::TensorIndex;
 
+namespace Rice::detail
+{
+  template<typename T>
+  struct Type<c10::complex<T>>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
+
+  template<typename T>
+  class To_Ruby<c10::complex<T>>
+  {
+  public:
+    VALUE convert(c10::complex<T> const& x)
+    {
+      return rb_dbl_complex_new(x.real(), x.imag());
+    }
+  };
+}
+
+template<typename T>
+Array flat_data(Tensor& tensor) {
+  Tensor view = tensor.reshape({tensor.numel()});
+
+  Array a;
+  for (int i = 0; i < tensor.numel(); i++) {
+    a.push(view[i].item().to<T>());
+  }
+  return a;
+}
+
 Class rb_cTensor;
 
 std::vector<TensorIndex> index_vector(Array a) {
@@ -23,19 +55,19 @@ std::vector<TensorIndex> index_vector(Array a) {
     obj = a[i];
 
     if (obj.is_instance_of(rb_cInteger)) {
-      indices.push_back(from_ruby<int64_t>(obj));
+      indices.push_back(Rice::detail::From_Ruby<int64_t>().convert(obj.value()));
     } else if (obj.is_instance_of(rb_cRange)) {
       torch::optional<int64_t> start_index = torch::nullopt;
       torch::optional<int64_t> stop_index = torch::nullopt;
 
       Object begin = obj.call("begin");
       if (!begin.is_nil()) {
-        start_index = from_ruby<int64_t>(begin);
+        start_index = Rice::detail::From_Ruby<int64_t>().convert(begin.value());
       }
 
       Object end = obj.call("end");
       if (!end.is_nil()) {
-        stop_index = from_ruby<int64_t>(end);
+        stop_index = Rice::detail::From_Ruby<int64_t>().convert(end.value());
       }
 
       Object exclude_end = obj.call("exclude_end?");
@@ -49,11 +81,11 @@ std::vector<TensorIndex> index_vector(Array a) {
 
       indices.push_back(torch::indexing::Slice(start_index, stop_index));
     } else if (obj.is_instance_of(rb_cTensor)) {
-      indices.push_back(from_ruby<Tensor>(obj));
+      indices.push_back(Rice::detail::From_Ruby<Tensor>().convert(obj.value()));
     } else if (obj.is_nil()) {
       indices.push_back(torch::indexing::None);
     } else if (obj == True || obj == False) {
-      indices.push_back(from_ruby<bool>(obj));
+      indices.push_back(Rice::detail::From_Ruby<bool>().convert(obj.value()));
     } else {
       throw Exception(rb_eArgError, "Unsupported index type: %s", rb_obj_classname(obj));
     }
@@ -68,7 +100,7 @@ std::vector<TensorIndex> index_vector(Array a) {
 static VALUE tensor__backward(int argc, VALUE* argv, VALUE self_)
 {
   HANDLE_TH_ERRORS
-  Tensor& self = from_ruby<Tensor&>(self_);
+  Tensor& self = Rice::detail::From_Ruby<Tensor&>().convert(self_);
   static RubyArgParser parser({
     "_backward(Tensor? gradient=None, bool? retain_graph=None, bool create_graph=False)"
   });
@@ -84,8 +116,8 @@ static VALUE tensor__backward(int argc, VALUE* argv, VALUE self_)
   END_HANDLE_TH_ERRORS
 }
 
-void init_tensor(Rice::Module& m) {
-  rb_cTensor = Rice::define_class_under<torch::Tensor>(m, "Tensor");
+void init_tensor(Rice::Module& m, Rice::Class& c, Rice::Class& rb_cTensorOptions) {
+  rb_cTensor = c;
   rb_cTensor.add_handler<torch::Error>(handle_error);
   add_tensor_functions(rb_cTensor);
   THPVariableClass = rb_cTensor.value();
@@ -102,18 +134,18 @@ void init_tensor(Rice::Module& m) {
     .define_method("requires_grad", &torch::Tensor::requires_grad)
     .define_method(
       "_size",
-      *[](Tensor& self, int64_t dim) {
+      [](Tensor& self, int64_t dim) {
         return self.size(dim);
       })
     .define_method(
       "_stride",
-      *[](Tensor& self, int64_t dim) {
+      [](Tensor& self, int64_t dim) {
         return self.stride(dim);
       })
     // in C++ for performance
     .define_method(
       "shape",
-      *[](Tensor& self) {
+      [](Tensor& self) {
         Array a;
         for (auto &size : self.sizes()) {
           a.push(size);
@@ -122,7 +154,7 @@ void init_tensor(Rice::Module& m) {
       })
     .define_method(
       "_strides",
-      *[](Tensor& self) {
+      [](Tensor& self) {
         Array a;
         for (auto &stride : self.strides()) {
           a.push(stride);
@@ -131,65 +163,65 @@ void init_tensor(Rice::Module& m) {
       })
     .define_method(
       "_index",
-      *[](Tensor& self, Array indices) {
+      [](Tensor& self, Array indices) {
         auto vec = index_vector(indices);
         return self.index(vec);
       })
     .define_method(
       "_index_put_custom",
-      *[](Tensor& self, Array indices, torch::Tensor& value) {
+      [](Tensor& self, Array indices, torch::Tensor& value) {
         auto vec = index_vector(indices);
         return self.index_put_(vec, value);
       })
     .define_method(
       "contiguous?",
-      *[](Tensor& self) {
+      [](Tensor& self) {
         return self.is_contiguous();
       })
     .define_method(
       "_requires_grad!",
-      *[](Tensor& self, bool requires_grad) {
+      [](Tensor& self, bool requires_grad) {
         return self.set_requires_grad(requires_grad);
       })
     .define_method(
       "grad",
-      *[](Tensor& self) {
+      [](Tensor& self) {
         auto grad = self.grad();
-        return grad.defined() ? to_ruby<torch::Tensor>(grad) : Nil;
+        return grad.defined() ? Object(Rice::detail::To_Ruby<torch::Tensor>().convert(grad)) : Nil;
       })
     .define_method(
       "grad=",
-      *[](Tensor& self, torch::Tensor& grad) {
+      [](Tensor& self, torch::Tensor& grad) {
         self.mutable_grad() = grad;
       })
     .define_method(
       "_dtype",
-      *[](Tensor& self) {
+      [](Tensor& self) {
         return (int) at::typeMetaToScalarType(self.dtype());
       })
     .define_method(
       "_type",
-      *[](Tensor& self, int dtype) {
+      [](Tensor& self, int dtype) {
         return self.toType((torch::ScalarType) dtype);
       })
     .define_method(
       "_layout",
-      *[](Tensor& self) {
+      [](Tensor& self) {
         std::stringstream s;
         s << self.layout();
         return s.str();
       })
     .define_method(
       "device",
-      *[](Tensor& self) {
+      [](Tensor& self) {
         std::stringstream s;
         s << self.device();
         return s.str();
       })
     .define_method(
       "_data_str",
-      *[](Tensor& self) {
-        Tensor tensor = self;
+      [](Tensor& self) {
+        auto tensor = self;
 
         // move to CPU to get data
         if (tensor.device().type() != torch::kCPU) {
@@ -207,14 +239,14 @@ void init_tensor(Rice::Module& m) {
     // for TorchVision
     .define_method(
       "_data_ptr",
-      *[](Tensor& self) {
+      [](Tensor& self) {
         return reinterpret_cast<uintptr_t>(self.data_ptr());
       })
     // TODO figure out a better way to do this
     .define_method(
       "_flat_data",
-      *[](Tensor& self) {
-        Tensor tensor = self;
+      [](Tensor& self) {
+        auto tensor = self;
 
         // move to CPU to get data
         if (tensor.device().type() != torch::kCPU) {
@@ -222,66 +254,47 @@ void init_tensor(Rice::Module& m) {
           tensor = tensor.to(device);
         }
 
-        Array a;
         auto dtype = tensor.dtype();
-
-        Tensor view = tensor.reshape({tensor.numel()});
-
-        // TODO DRY if someone knows C++
         if (dtype == torch::kByte) {
-          for (int i = 0; i < tensor.numel(); i++) {
-            a.push(view[i].item().to<uint8_t>());
-          }
+          return flat_data<uint8_t>(tensor);
         } else if (dtype == torch::kChar) {
-          for (int i = 0; i < tensor.numel(); i++) {
-            a.push(to_ruby<int>(view[i].item().to<int8_t>()));
-          }
+          return flat_data<int8_t>(tensor);
         } else if (dtype == torch::kShort) {
-          for (int i = 0; i < tensor.numel(); i++) {
-            a.push(view[i].item().to<int16_t>());
-          }
+          return flat_data<int16_t>(tensor);
         } else if (dtype == torch::kInt) {
-          for (int i = 0; i < tensor.numel(); i++) {
-            a.push(view[i].item().to<int32_t>());
-          }
+          return flat_data<int32_t>(tensor);
         } else if (dtype == torch::kLong) {
-          for (int i = 0; i < tensor.numel(); i++) {
-            a.push(view[i].item().to<int64_t>());
-          }
+          return flat_data<int64_t>(tensor);
         } else if (dtype == torch::kFloat) {
-          for (int i = 0; i < tensor.numel(); i++) {
-            a.push(view[i].item().to<float>());
-          }
+          return flat_data<float>(tensor);
         } else if (dtype == torch::kDouble) {
-          for (int i = 0; i < tensor.numel(); i++) {
-            a.push(view[i].item().to<double>());
-          }
+          return flat_data<double>(tensor);
         } else if (dtype == torch::kBool) {
-          for (int i = 0; i < tensor.numel(); i++) {
-            a.push(view[i].item().to<bool>() ? True : False);
-          }
+          return flat_data<bool>(tensor);
+        } else if (dtype == torch::kComplexFloat) {
+          return flat_data<c10::complex<float>>(tensor);
+        } else if (dtype == torch::kComplexDouble) {
+          return flat_data<c10::complex<double>>(tensor);
         } else {
           throw std::runtime_error("Unsupported type");
         }
-        return a;
       })
     .define_method(
       "_to",
-      *[](Tensor& self, torch::Device device, int dtype, bool non_blocking, bool copy) {
+      [](Tensor& self, torch::Device device, int dtype, bool non_blocking, bool copy) {
         return self.to(device, (torch::ScalarType) dtype, non_blocking, copy);
       });
 
-  Rice::define_class_under<torch::TensorOptions>(m, "TensorOptions")
+  rb_cTensorOptions
     .add_handler<torch::Error>(handle_error)
-    .define_constructor(Rice::Constructor<torch::TensorOptions>())
     .define_method(
       "dtype",
-      *[](torch::TensorOptions& self, int dtype) {
+      [](torch::TensorOptions& self, int dtype) {
         return self.dtype((torch::ScalarType) dtype);
       })
     .define_method(
       "layout",
-      *[](torch::TensorOptions& self, const std::string& layout) {
+      [](torch::TensorOptions& self, const std::string& layout) {
         torch::Layout l;
         if (layout == "strided") {
           l = torch::kStrided;
@@ -295,13 +308,13 @@ void init_tensor(Rice::Module& m) {
       })
     .define_method(
       "device",
-      *[](torch::TensorOptions& self, const std::string& device) {
+      [](torch::TensorOptions& self, const std::string& device) {
         torch::Device d(device);
         return self.device(d);
       })
     .define_method(
       "requires_grad",
-      *[](torch::TensorOptions& self, bool requires_grad) {
+      [](torch::TensorOptions& self, bool requires_grad) {
         return self.requires_grad(requires_grad);
       });
 }