torch-rb 0.5.0 → 0.7.0

data/ext/torch/extconf.rb

@@ -1,8 +1,6 @@
 require "mkmf-rice"
 
-abort "Missing stdc++" unless have_library("stdc++")
-
-$CXXFLAGS += " -std=c++14"
+$CXXFLAGS += " -std=c++17 $(optflags)"
 
 # change to 0 for Linux pre-cxx11 ABI version
 $CXXFLAGS += " -D_GLIBCXX_USE_CXX11_ABI=1"
@@ -11,7 +9,6 @@ apple_clang = RbConfig::CONFIG["CC_VERSION_MESSAGE"] =~ /apple clang/i
 
 # check omp first
 if have_library("omp") || have_library("gomp")
-  $CXXFLAGS += " -DAT_PARALLEL_OPENMP=1"
   $CXXFLAGS += " -Xclang" if apple_clang
   $CXXFLAGS += " -fopenmp"
 end

data/ext/torch/ivalue.cpp

@@ -0,0 +1,132 @@
+#include <torch/torch.h>
+
+#include <rice/rice.hpp>
+
+#include "utils.h"
+
+void init_ivalue(Rice::Module& m, Rice::Class& rb_cIValue) {
+  // https://pytorch.org/cppdocs/api/structc10_1_1_i_value.html
+  rb_cIValue
+    .add_handler<torch::Error>(handle_error)
+    .define_method("bool?", &torch::IValue::isBool)
+    .define_method("bool_list?", &torch::IValue::isBoolList)
+    .define_method("capsule?", &torch::IValue::isCapsule)
+    .define_method("custom_class?", &torch::IValue::isCustomClass)
+    .define_method("device?", &torch::IValue::isDevice)
+    .define_method("double?", &torch::IValue::isDouble)
+    .define_method("double_list?", &torch::IValue::isDoubleList)
+    .define_method("future?", &torch::IValue::isFuture)
+    // .define_method("generator?", &torch::IValue::isGenerator)
+    .define_method("generic_dict?", &torch::IValue::isGenericDict)
+    .define_method("list?", &torch::IValue::isList)
+    .define_method("int?", &torch::IValue::isInt)
+    .define_method("int_list?", &torch::IValue::isIntList)
+    .define_method("module?", &torch::IValue::isModule)
+    .define_method("none?", &torch::IValue::isNone)
+    .define_method("object?", &torch::IValue::isObject)
+    .define_method("ptr_type?", &torch::IValue::isPtrType)
+    .define_method("py_object?", &torch::IValue::isPyObject)
+    .define_method("r_ref?", &torch::IValue::isRRef)
+    .define_method("scalar?", &torch::IValue::isScalar)
+    .define_method("string?", &torch::IValue::isString)
+    .define_method("tensor?", &torch::IValue::isTensor)
+    .define_method("tensor_list?", &torch::IValue::isTensorList)
+    .define_method("tuple?", &torch::IValue::isTuple)
+    .define_method(
+      "to_bool",
+      [](torch::IValue& self) {
+        return self.toBool();
+      })
+    .define_method(
+      "to_double",
+      [](torch::IValue& self) {
+        return self.toDouble();
+      })
+    .define_method(
+      "to_int",
+      [](torch::IValue& self) {
+        return self.toInt();
+      })
+    .define_method(
+      "to_list",
+      [](torch::IValue& self) {
+        auto list = self.toListRef();
+        Rice::Array obj;
+        for (auto& elem : list) {
+          auto v = torch::IValue{elem};
+          obj.push(Rice::Object(Rice::detail::To_Ruby<torch::IValue>().convert(v)));
+        }
+        return obj;
+      })
+    .define_method(
+      "to_string_ref",
+      [](torch::IValue& self) {
+        return self.toStringRef();
+      })
+    .define_method(
+      "to_tensor",
+      [](torch::IValue& self) {
+        return self.toTensor();
+      })
+    .define_method(
+      "to_generic_dict",
+      [](torch::IValue& self) {
+        auto dict = self.toGenericDict();
+        Rice::Hash obj;
+        for (auto& pair : dict) {
+          auto k = torch::IValue{pair.key()};
+          auto v = torch::IValue{pair.value()};
+          obj[Rice::Object(Rice::detail::To_Ruby<torch::IValue>().convert(k))] = Rice::Object(Rice::detail::To_Ruby<torch::IValue>().convert(v));
+        }
+        return obj;
+      })
+    .define_singleton_function(
+      "from_tensor",
+      [](torch::Tensor& v) {
+        return torch::IValue(v);
+      })
+    // TODO create specialized list types?
+    .define_singleton_function(
+      "from_list",
+      [](Rice::Array obj) {
+        c10::impl::GenericList list(c10::AnyType::get());
+        for (auto entry : obj) {
+          list.push_back(Rice::detail::From_Ruby<torch::IValue>().convert(entry.value()));
+        }
+        return torch::IValue(list);
+      })
+    .define_singleton_function(
+      "from_string",
+      [](Rice::String v) {
+        return torch::IValue(v.str());
+      })
+    .define_singleton_function(
+      "from_int",
+      [](int64_t v) {
+        return torch::IValue(v);
+      })
+    .define_singleton_function(
+      "from_double",
+      [](double v) {
+        return torch::IValue(v);
+      })
+    .define_singleton_function(
+      "from_bool",
+      [](bool v) {
+        return torch::IValue(v);
+      })
+    // see https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/python/pybind_utils.h
+    // createGenericDict and toIValue
+    .define_singleton_function(
+      "from_dict",
+      [](Rice::Hash obj) {
+        auto key_type = c10::AnyType::get();
+        auto value_type = c10::AnyType::get();
+        c10::impl::GenericDict elems(key_type, value_type);
+        elems.reserve(obj.size());
+        for (auto entry : obj) {
+          elems.insert(Rice::detail::From_Ruby<torch::IValue>().convert(entry.first), Rice::detail::From_Ruby<torch::IValue>().convert((Rice::Object) entry.second));
+        }
+        return torch::IValue(std::move(elems));
+      });
+}

data/ext/torch/nn.cpp

@@ -0,0 +1,114 @@
+#include <torch/torch.h>
+
+#include <rice/rice.hpp>
+
+#include "nn_functions.h"
+#include "templates.h"
+#include "utils.h"
+
+// need to make a distinction between parameters and tensors
+class Parameter: public torch::autograd::Variable {
+  public:
+    Parameter(Tensor&& t) : torch::autograd::Variable(t) { }
+};
+
+void init_nn(Rice::Module& m) {
+  auto rb_mNN = Rice::define_module_under(m, "NN");
+  rb_mNN.add_handler<torch::Error>(handle_error);
+  add_nn_functions(rb_mNN);
+
+  Rice::define_module_under(rb_mNN, "Init")
+    .add_handler<torch::Error>(handle_error)
+    .define_singleton_function(
+      "_calculate_gain",
+      [](NonlinearityType nonlinearity, double param) {
+        return torch::nn::init::calculate_gain(nonlinearity, param);
+      })
+    .define_singleton_function(
+      "_uniform!",
+      [](Tensor tensor, double low, double high) {
+        return torch::nn::init::uniform_(tensor, low, high);
+      })
+    .define_singleton_function(
+      "_normal!",
+      [](Tensor tensor, double mean, double std) {
+        return torch::nn::init::normal_(tensor, mean, std);
+      })
+    .define_singleton_function(
+      "_constant!",
+      [](Tensor tensor, Scalar value) {
+        return torch::nn::init::constant_(tensor, value);
+      })
+    .define_singleton_function(
+      "_ones!",
+      [](Tensor tensor) {
+        return torch::nn::init::ones_(tensor);
+      })
+    .define_singleton_function(
+      "_zeros!",
+      [](Tensor tensor) {
+        return torch::nn::init::zeros_(tensor);
+      })
+    .define_singleton_function(
+      "_eye!",
+      [](Tensor tensor) {
+        return torch::nn::init::eye_(tensor);
+      })
+    .define_singleton_function(
+      "_dirac!",
+      [](Tensor tensor) {
+        return torch::nn::init::dirac_(tensor);
+      })
+    .define_singleton_function(
+      "_xavier_uniform!",
+      [](Tensor tensor, double gain) {
+        return torch::nn::init::xavier_uniform_(tensor, gain);
+      })
+    .define_singleton_function(
+      "_xavier_normal!",
+      [](Tensor tensor, double gain) {
+        return torch::nn::init::xavier_normal_(tensor, gain);
+      })
+    .define_singleton_function(
+      "_kaiming_uniform!",
+      [](Tensor tensor, double a, FanModeType mode, NonlinearityType nonlinearity) {
+        return torch::nn::init::kaiming_uniform_(tensor, a, mode, nonlinearity);
+      })
+    .define_singleton_function(
+      "_kaiming_normal!",
+      [](Tensor tensor, double a, FanModeType mode, NonlinearityType nonlinearity) {
+        return torch::nn::init::kaiming_normal_(tensor, a, mode, nonlinearity);
+      })
+    .define_singleton_function(
+      "_orthogonal!",
+      [](Tensor tensor, double gain) {
+        return torch::nn::init::orthogonal_(tensor, gain);
+      })
+    .define_singleton_function(
+      "_sparse!",
+      [](Tensor tensor, double sparsity, double std) {
+        return torch::nn::init::sparse_(tensor, sparsity, std);
+      });
+
+  Rice::define_class_under<Parameter, torch::Tensor>(rb_mNN, "Parameter")
+    .add_handler<torch::Error>(handle_error)
+    .define_method(
+      "grad",
+      [](Parameter& self) {
+        auto grad = self.grad();
+        return grad.defined() ? Object(Rice::detail::To_Ruby<torch::Tensor>().convert(grad)) : Nil;
+      })
+    .define_method(
+      "grad=",
+      [](Parameter& self, torch::Tensor& grad) {
+        self.mutable_grad() = grad;
+      })
+    .define_singleton_function(
+      "_make_subclass",
+      [](Tensor& rd, bool requires_grad) {
+        auto data = rd.detach();
+        data.unsafeGetTensorImpl()->set_allow_tensor_metadata_change(true);
+        auto var = data.set_requires_grad(requires_grad);
+        return Parameter(std::move(var));
+      });
+}

data/ext/torch/nn_functions.h

@@ -3,4 +3,4 @@
 
 #pragma once
 
-void add_nn_functions(Module m);
+void add_nn_functions(Rice::Module& m);

data/ext/torch/random.cpp

@@ -0,0 +1,22 @@
+#include <torch/torch.h>
+
+#include <rice/rice.hpp>
+
+#include "utils.h"
+
+void init_random(Rice::Module& m) {
+  Rice::define_module_under(m, "Random")
+    .add_handler<torch::Error>(handle_error)
+    .define_singleton_function(
+      "initial_seed",
+      []() {
+        return at::detail::getDefaultCPUGenerator().current_seed();
+      })
+    .define_singleton_function(
+      "seed",
+      []() {
+        // TODO set for CUDA when available
+        auto generator = at::detail::getDefaultCPUGenerator();
+        return generator.seed();
+      });
+}

data/ext/torch/ruby_arg_parser.cpp

@@ -137,7 +137,7 @@ auto FunctionParameter::check(VALUE obj, int argnum) -> bool
         return true;
       }
       if (THPVariable_Check(obj)) {
-        auto var = from_ruby<torch::Tensor>(obj);
+        auto var = Rice::detail::From_Ruby<torch::Tensor>().convert(obj);
         return !var.requires_grad() && var.dim() == 0;
       }
       return false;
@@ -147,7 +147,7 @@ auto FunctionParameter::check(VALUE obj, int argnum) -> bool
         return true;
       }
       if (THPVariable_Check(obj)) {
-        auto var = from_ruby<torch::Tensor>(obj);
+        auto var = Rice::detail::From_Ruby<torch::Tensor>().convert(obj);
         return at::isIntegralType(var.scalar_type(), /*includeBool=*/false) && !var.requires_grad() && var.dim() == 0;
       }
       return false;
@@ -487,7 +487,7 @@ static void extra_kwargs(FunctionSignature& signature, VALUE kwargs, ssize_t num
 
 VALUE missing = Qundef;
 
-bool FunctionSignature::parse(VALUE self, VALUE args, VALUE kwargs, std::vector<VALUE> &dst,  // NOLINT
+bool FunctionSignature::parse(VALUE self, VALUE args, VALUE kwargs, VALUE dst[],  // NOLINT
                               bool raise_exception) {
   auto nargs = NIL_P(args) ? 0 : RARRAY_LEN(args);
   ssize_t remaining_kwargs = NIL_P(kwargs) ? 0 :  RHASH_SIZE(kwargs);

data/ext/torch/ruby_arg_parser.h

@@ -2,8 +2,10 @@
 
 #pragma once
 
+#include <sstream>
+
 #include <torch/torch.h>
-#include <rice/Exception.hpp>
+#include <rice/rice.hpp>
 
 #include "templates.h"
 #include "utils.h"
@@ -46,7 +48,7 @@ struct FunctionParameter {
 struct FunctionSignature {
   explicit FunctionSignature(const std::string& fmt, int index);
 
-  bool parse(VALUE self, VALUE args, VALUE kwargs, std::vector<VALUE>& dst, bool raise_exception);
+  bool parse(VALUE self, VALUE args, VALUE kwargs, VALUE dst[], bool raise_exception);
 
   std::string toString() const;
 
@@ -63,13 +65,13 @@ struct FunctionSignature {
 };
 
 struct RubyArgs {
-  RubyArgs(const FunctionSignature& signature, std::vector<VALUE> &args)
+  RubyArgs(const FunctionSignature& signature, VALUE* args)
     : signature(signature)
     , args(args)
     , idx(signature.index) {}
 
   const FunctionSignature& signature;
-  std::vector<VALUE> args;
+  VALUE* args;
   int idx;
 
   inline at::Tensor tensor(int i);
@@ -119,7 +121,7 @@ struct RubyArgs {
 };
 
 inline at::Tensor RubyArgs::tensor(int i) {
-  return from_ruby<torch::Tensor>(args[i]);
+  return Rice::detail::From_Ruby<torch::Tensor>().convert(args[i]);
 }
 
 inline OptionalTensor RubyArgs::optionalTensor(int i) {
@@ -129,12 +131,12 @@ inline OptionalTensor RubyArgs::optionalTensor(int i) {
 
 inline at::Scalar RubyArgs::scalar(int i) {
   if (NIL_P(args[i])) return signature.params[i].default_scalar;
-  return from_ruby<torch::Scalar>(args[i]);
+  return Rice::detail::From_Ruby<torch::Scalar>().convert(args[i]);
 }
 
 inline std::vector<at::Tensor> RubyArgs::tensorlist(int i) {
   if (NIL_P(args[i])) return std::vector<at::Tensor>();
-  return from_ruby<std::vector<Tensor>>(args[i]);
+  return Rice::detail::From_Ruby<std::vector<Tensor>>().convert(args[i]);
 }
 
 template<int N>
@@ -149,7 +151,7 @@ inline std::array<at::Tensor, N> RubyArgs::tensorlist_n(int i) {
   }
   for (int idx = 0; idx < size; idx++) {
     VALUE obj = rb_ary_entry(arg, idx);
-    res[idx] = from_ruby<Tensor>(obj);
+    res[idx] = Rice::detail::From_Ruby<Tensor>().convert(obj);
   }
   return res;
 }
@@ -168,7 +170,7 @@ inline std::vector<int64_t> RubyArgs::intlist(int i) {
   for (idx = 0; idx < size; idx++) {
     VALUE obj = rb_ary_entry(arg, idx);
     if (FIXNUM_P(obj)) {
-      res[idx] = from_ruby<int64_t>(obj);
+      res[idx] = Rice::detail::From_Ruby<int64_t>().convert(obj);
     } else {
       rb_raise(rb_eArgError, "%s(): argument '%s' must be %s, but found element of type %s at pos %d",
           signature.name.c_str(), signature.params[i].name.c_str(),
@@ -208,8 +210,13 @@ inline ScalarType RubyArgs::scalartype(int i) {
     {ID2SYM(rb_intern("double")), ScalarType::Double},
     {ID2SYM(rb_intern("float64")), ScalarType::Double},
     {ID2SYM(rb_intern("complex_half")), ScalarType::ComplexHalf},
+    {ID2SYM(rb_intern("complex32")), ScalarType::ComplexHalf},
     {ID2SYM(rb_intern("complex_float")), ScalarType::ComplexFloat},
+    {ID2SYM(rb_intern("cfloat")), ScalarType::ComplexFloat},
+    {ID2SYM(rb_intern("complex64")), ScalarType::ComplexFloat},
     {ID2SYM(rb_intern("complex_double")), ScalarType::ComplexDouble},
+    {ID2SYM(rb_intern("cdouble")), ScalarType::ComplexDouble},
+    {ID2SYM(rb_intern("complex128")), ScalarType::ComplexDouble},
     {ID2SYM(rb_intern("bool")), ScalarType::Bool},
     {ID2SYM(rb_intern("qint8")), ScalarType::QInt8},
     {ID2SYM(rb_intern("quint8")), ScalarType::QUInt8},
@@ -258,7 +265,7 @@ inline c10::OptionalArray<double> RubyArgs::doublelistOptional(int i) {
   for (idx = 0; idx < size; idx++) {
     VALUE obj = rb_ary_entry(arg, idx);
     if (FIXNUM_P(obj) || RB_FLOAT_TYPE_P(obj)) {
-      res[idx] = from_ruby<double>(obj);
+      res[idx] = Rice::detail::From_Ruby<double>().convert(obj);
     } else {
       rb_raise(rb_eArgError, "%s(): argument '%s' must be %s, but found element of type %s at pos %d",
           signature.name.c_str(), signature.params[i].name.c_str(),
@@ -301,22 +308,22 @@ inline c10::optional<at::MemoryFormat> RubyArgs::memoryformatOptional(int i) {
 }
 
 inline std::string RubyArgs::string(int i) {
-  return from_ruby<std::string>(args[i]);
+  return Rice::detail::From_Ruby<std::string>().convert(args[i]);
 }
 
 inline c10::optional<std::string> RubyArgs::stringOptional(int i) {
   if (!args[i]) return c10::nullopt;
-  return from_ruby<std::string>(args[i]);
+  return Rice::detail::From_Ruby<std::string>().convert(args[i]);
 }
 
 inline int64_t RubyArgs::toInt64(int i) {
   if (NIL_P(args[i])) return signature.params[i].default_int;
-  return from_ruby<int64_t>(args[i]);
+  return Rice::detail::From_Ruby<int64_t>().convert(args[i]);
 }
 
 inline double RubyArgs::toDouble(int i) {
   if (NIL_P(args[i])) return signature.params[i].default_double;
-  return from_ruby<double>(args[i]);
+  return Rice::detail::From_Ruby<double>().convert(args[i]);
 }
 
 inline bool RubyArgs::toBool(int i) {
@@ -328,6 +335,12 @@ inline bool RubyArgs::isNone(int i) {
   return NIL_P(args[i]);
 }
 
+template<int N>
+struct ParsedArgs {
+  ParsedArgs() : args() { }
+  VALUE args[N];
+};
+
 struct RubyArgParser {
   std::vector<FunctionSignature> signatures_;
   std::string function_name;
@@ -356,7 +369,15 @@ struct RubyArgParser {
         });
     }
 
-    RubyArgs parse(VALUE self, int argc, VALUE* argv, std::vector<VALUE> &parsed_args) {
+    template<int N>
+    inline RubyArgs parse(VALUE self, int argc, VALUE* argv, ParsedArgs<N> &dst) {
+      if (N < max_args) {
+        rb_raise(rb_eArgError, "RubyArgParser: dst ParsedArgs buffer does not have enough capacity, expected %d (got %d)", (int)max_args, N);
+      }
+      return raw_parse(self, argc, argv, dst.args);
+    }
+
+    inline RubyArgs raw_parse(VALUE self, int argc, VALUE* argv, VALUE parsed_args[]) {
       VALUE args, kwargs;
       rb_scan_args(argc, argv, "*:", &args, &kwargs);
 
@@ -378,7 +399,7 @@ struct RubyArgParser {
       rb_raise(rb_eArgError, "No matching signatures");
     }
 
-    void print_error(VALUE self, VALUE args, VALUE kwargs, std::vector<VALUE>& parsed_args) {
+    void print_error(VALUE self, VALUE args, VALUE kwargs, VALUE parsed_args[]) {
       ssize_t num_args = (NIL_P(args) ? 0 : RARRAY_LEN(args)) + (NIL_P(kwargs) ? 0 : RHASH_SIZE(kwargs));
       std::vector<int> plausible_idxs;
       ssize_t i = 0;