torch-rb 0.4.2 → 0.6.0

data/ext/torch/extconf.rb

@@ -11,7 +11,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,134 @@
+#include <torch/torch.h>
+
+#include <rice/Array.hpp>
+#include <rice/Constructor.hpp>
+#include <rice/Hash.hpp>
+#include <rice/Module.hpp>
+#include <rice/String.hpp>
+
+#include "utils.h"
+
+void init_ivalue(Rice::Module& m) {
+  // https://pytorch.org/cppdocs/api/structc10_1_1_i_value.html
+  Rice::define_class_under<torch::IValue>(m, "IValue")
+    .add_handler<torch::Error>(handle_error)
+    .define_constructor(Rice::Constructor<torch::IValue>())
+    .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) {
+          obj.push(to_ruby<torch::IValue>(torch::IValue{elem}));
+        }
+        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) {
+          obj[to_ruby<torch::IValue>(torch::IValue{pair.key()})] = to_ruby<torch::IValue>(torch::IValue{pair.value()});
+        }
+        return obj;
+      })
+    .define_singleton_method(
+      "from_tensor",
+      *[](torch::Tensor& v) {
+        return torch::IValue(v);
+      })
+    // TODO create specialized list types?
+    .define_singleton_method(
+      "from_list",
+      *[](Rice::Array obj) {
+        c10::impl::GenericList list(c10::AnyType::get());
+        for (auto entry : obj) {
+          list.push_back(from_ruby<torch::IValue>(entry));
+        }
+        return torch::IValue(list);
+      })
+    .define_singleton_method(
+      "from_string",
+      *[](Rice::String v) {
+        return torch::IValue(v.str());
+      })
+    .define_singleton_method(
+      "from_int",
+      *[](int64_t v) {
+        return torch::IValue(v);
+      })
+    .define_singleton_method(
+      "from_double",
+      *[](double v) {
+        return torch::IValue(v);
+      })
+    .define_singleton_method(
+      "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_method(
+      "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(from_ruby<torch::IValue>(entry.first), from_ruby<torch::IValue>((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/Module.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_method(
+      "_calculate_gain",
+      *[](NonlinearityType nonlinearity, double param) {
+        return torch::nn::init::calculate_gain(nonlinearity, param);
+      })
+    .define_singleton_method(
+      "_uniform!",
+      *[](Tensor tensor, double low, double high) {
+        return torch::nn::init::uniform_(tensor, low, high);
+      })
+    .define_singleton_method(
+      "_normal!",
+      *[](Tensor tensor, double mean, double std) {
+        return torch::nn::init::normal_(tensor, mean, std);
+      })
+    .define_singleton_method(
+      "_constant!",
+      *[](Tensor tensor, Scalar value) {
+        return torch::nn::init::constant_(tensor, value);
+      })
+    .define_singleton_method(
+      "_ones!",
+      *[](Tensor tensor) {
+        return torch::nn::init::ones_(tensor);
+      })
+    .define_singleton_method(
+      "_zeros!",
+      *[](Tensor tensor) {
+        return torch::nn::init::zeros_(tensor);
+      })
+    .define_singleton_method(
+      "_eye!",
+      *[](Tensor tensor) {
+        return torch::nn::init::eye_(tensor);
+      })
+    .define_singleton_method(
+      "_dirac!",
+      *[](Tensor tensor) {
+        return torch::nn::init::dirac_(tensor);
+      })
+    .define_singleton_method(
+      "_xavier_uniform!",
+      *[](Tensor tensor, double gain) {
+        return torch::nn::init::xavier_uniform_(tensor, gain);
+      })
+    .define_singleton_method(
+      "_xavier_normal!",
+      *[](Tensor tensor, double gain) {
+        return torch::nn::init::xavier_normal_(tensor, gain);
+      })
+    .define_singleton_method(
+      "_kaiming_uniform!",
+      *[](Tensor tensor, double a, FanModeType mode, NonlinearityType nonlinearity) {
+        return torch::nn::init::kaiming_uniform_(tensor, a, mode, nonlinearity);
+      })
+    .define_singleton_method(
+      "_kaiming_normal!",
+      *[](Tensor tensor, double a, FanModeType mode, NonlinearityType nonlinearity) {
+        return torch::nn::init::kaiming_normal_(tensor, a, mode, nonlinearity);
+      })
+    .define_singleton_method(
+      "_orthogonal!",
+      *[](Tensor tensor, double gain) {
+        return torch::nn::init::orthogonal_(tensor, gain);
+      })
+    .define_singleton_method(
+      "_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() ? to_ruby<torch::Tensor>(grad) : Nil;
+      })
+    .define_method(
+      "grad=",
+      *[](Parameter& self, torch::Tensor& grad) {
+        self.mutable_grad() = grad;
+      })
+    .define_singleton_method(
+      "_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/Module.hpp>
+
+#include "utils.h"
+
+void init_random(Rice::Module& m) {
+  Rice::define_module_under(m, "Random")
+    .add_handler<torch::Error>(handle_error)
+    .define_singleton_method(
+      "initial_seed",
+      *[]() {
+        return at::detail::getDefaultCPUGenerator().current_seed();
+      })
+    .define_singleton_method(
+      "seed",
+      *[]() {
+        // TODO set for CUDA when available
+        auto generator = at::detail::getDefaultCPUGenerator();
+        return generator.seed();
+      });
+}

data/ext/torch/ruby_arg_parser.cpp

@@ -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,6 +2,8 @@
 
 #pragma once
 
+#include <sstream>
+
 #include <torch/torch.h>
 #include <rice/Exception.hpp>
 
@@ -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);
@@ -91,7 +93,7 @@ struct RubyArgs {
   inline c10::optional<int64_t> toInt64Optional(int i);
   inline c10::optional<bool> toBoolOptional(int i);
   inline c10::optional<double> toDoubleOptional(int i);
-  // inline c10::OptionalArray<double> doublelistOptional(int i);
+  inline c10::OptionalArray<double> doublelistOptional(int i);
   // inline at::Layout layout(int i);
   // inline at::Layout layoutWithDefault(int i, at::Layout default_layout);
   inline c10::optional<at::Layout> layoutOptional(int i);
@@ -105,7 +107,7 @@ struct RubyArgs {
   inline c10::optional<at::MemoryFormat> memoryformatOptional(int i);
   // inline at::QScheme toQScheme(int i);
   inline std::string string(int i);
-  // inline c10::optional<std::string> stringOptional(int i);
+  inline c10::optional<std::string> stringOptional(int i);
   // inline PyObject* pyobject(int i);
   inline int64_t toInt64(int i);
   // inline int64_t toInt64WithDefault(int i, int64_t default_int);
@@ -249,6 +251,25 @@ inline c10::optional<double> RubyArgs::toDoubleOptional(int i) {
   return toDouble(i);
 }
 
+inline c10::OptionalArray<double> RubyArgs::doublelistOptional(int i) {
+  if (NIL_P(args[i])) return {};
+
+  VALUE arg = args[i];
+  auto size = RARRAY_LEN(arg);
+  std::vector<double> res(size);
+  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);
+    } 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(),
+          signature.params[i].type_name().c_str(), rb_obj_classname(obj), idx + 1);
+    }
+  }
+  return res;
+}
+
 inline c10::optional<at::Layout> RubyArgs::layoutOptional(int i) {
   if (NIL_P(args[i])) return c10::nullopt;
 
@@ -285,6 +306,11 @@ inline std::string RubyArgs::string(int i) {
   return from_ruby<std::string>(args[i]);
 }
 
+inline c10::optional<std::string> RubyArgs::stringOptional(int i) {
+  if (!args[i]) return c10::nullopt;
+  return from_ruby<std::string>(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]);
@@ -304,6 +330,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;
@@ -332,7 +364,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);
 
@@ -354,7 +394,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;

data/ext/torch/templates.h

@@ -19,6 +19,7 @@ using torch::TensorOptions;
 using torch::Layout;
 using torch::MemoryFormat;
 using torch::IntArrayRef;
+using torch::ArrayRef;
 using torch::TensorList;
 using torch::Storage;
 
@@ -43,7 +44,7 @@ std::vector<int64_t> from_ruby<std::vector<int64_t>>(Object x)
 {
   Array a = Array(x);
   std::vector<int64_t> vec(a.size());
-  for (size_t i = 0; i < a.size(); i++) {
+  for (long i = 0; i < a.size(); i++) {
     vec[i] = from_ruby<int64_t>(a[i]);
   }
   return vec;
@@ -55,7 +56,7 @@ std::vector<Tensor> from_ruby<std::vector<Tensor>>(Object x)
 {
   Array a = Array(x);
   std::vector<Tensor> vec(a.size());
-  for (size_t i = 0; i < a.size(); i++) {
+  for (long i = 0; i < a.size(); i++) {
     vec[i] = from_ruby<Tensor>(a[i]);
   }
   return vec;