torch-rb 0.5.2 → 0.8.1

data/ext/torch/tensor_functions.h

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

data/ext/torch/torch.cpp

@@ -0,0 +1,95 @@
+#include <torch/torch.h>
+
+#include <rice/rice.hpp>
+
+#include "torch_functions.h"
+#include "templates.h"
+#include "utils.h"
+
+void init_torch(Rice::Module& m) {
+  m.add_handler<torch::Error>(handle_error);
+  add_torch_functions(m);
+  m.define_singleton_function(
+      "grad_enabled?",
+      []() {
+        return torch::GradMode::is_enabled();
+      })
+    .define_singleton_function(
+      "_set_grad_enabled",
+      [](bool enabled) {
+        torch::GradMode::set_enabled(enabled);
+      })
+    .define_singleton_function(
+      "manual_seed",
+      [](uint64_t seed) {
+        return torch::manual_seed(seed);
+      })
+    // config
+    .define_singleton_function(
+      "show_config",
+      [] {
+        return torch::show_config();
+      })
+    .define_singleton_function(
+      "parallel_info",
+      [] {
+        return torch::get_parallel_info();
+      })
+    // begin operations
+    .define_singleton_function(
+      "_save",
+      [](const torch::IValue &value) {
+        auto v = torch::pickle_save(value);
+        std::string str(v.begin(), v.end());
+        return str;
+      })
+    .define_singleton_function(
+      "_load",
+      [](const std::string &s) {
+        std::vector<char> v;
+        std::copy(s.begin(), s.end(), std::back_inserter(v));
+        // https://github.com/pytorch/pytorch/issues/20356#issuecomment-567663701
+        return torch::pickle_load(v);
+      })
+    .define_singleton_function(
+      "_from_blob",
+      [](Rice::String s, std::vector<int64_t> size, const torch::TensorOptions &options) {
+        void *data = const_cast<char *>(s.c_str());
+        return torch::from_blob(data, size, options);
+      })
+    .define_singleton_function(
+      "_tensor",
+      [](Rice::Array a, std::vector<int64_t> size, const torch::TensorOptions &options) {
+        auto dtype = options.dtype();
+        torch::Tensor t;
+        if (dtype == torch::kBool) {
+          std::vector<uint8_t> vec;
+          for (long i = 0; i < a.size(); i++) {
+            vec.push_back(Rice::detail::From_Ruby<bool>().convert(a[i].value()));
+          }
+          t = torch::tensor(vec, options);
+        } else if (dtype == torch::kComplexFloat || dtype == torch::kComplexDouble) {
+          // TODO use template
+          std::vector<c10::complex<double>> vec;
+          Object obj;
+          for (long i = 0; i < a.size(); i++) {
+            obj = a[i];
+            vec.push_back(c10::complex<double>(Rice::detail::From_Ruby<double>().convert(obj.call("real").value()), Rice::detail::From_Ruby<double>().convert(obj.call("imag").value())));
+          }
+          t = torch::tensor(vec, options);
+        } else {
+          std::vector<float> vec;
+          for (long i = 0; i < a.size(); i++) {
+            vec.push_back(Rice::detail::From_Ruby<float>().convert(a[i].value()));
+          }
+          // hack for requires_grad error
+          if (options.requires_grad()) {
+            t = torch::tensor(vec, options.requires_grad(c10::nullopt));
+            t.set_requires_grad(true);
+          } else {
+            t = torch::tensor(vec, options);
+          }
+        }
+        return t.reshape(size);
+      });
+}

data/ext/torch/torch_functions.h

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

data/ext/torch/utils.h

@@ -1,13 +1,19 @@
 #pragma once
 
-#include <rice/Symbol.hpp>
+#include <rice/rice.hpp>
+#include <rice/stl.hpp>
+
+// TODO find better place
+inline void handle_error(torch::Error const & ex) {
+  throw Rice::Exception(rb_eRuntimeError, ex.what_without_backtrace());
+}
 
 // keep THP prefix for now to make it easier to compare code
 
 extern VALUE THPVariableClass;
 
 inline VALUE THPUtils_internSymbol(const std::string& str) {
-  return Symbol(str);
+  return Rice::Symbol(str);
 }
 
 inline std::string THPUtils_unpackSymbol(VALUE obj) {

data/ext/torch/wrap_outputs.h

@@ -1,99 +1,106 @@
 #pragma once
 
 #include <torch/torch.h>
-#include <rice/Object.hpp>
+#include <rice/rice.hpp>
 
-inline Object wrap(bool x) {
-  return to_ruby<bool>(x);
+inline VALUE wrap(bool x) {
+  return Rice::detail::To_Ruby<bool>().convert(x);
 }
 
-inline Object wrap(int64_t x) {
-  return to_ruby<int64_t>(x);
+inline VALUE wrap(int64_t x) {
+  return Rice::detail::To_Ruby<int64_t>().convert(x);
 }
 
-inline Object wrap(double x) {
-  return to_ruby<double>(x);
+inline VALUE wrap(double x) {
+  return Rice::detail::To_Ruby<double>().convert(x);
 }
 
-inline Object wrap(torch::Tensor x) {
-  return to_ruby<torch::Tensor>(x);
+inline VALUE wrap(torch::Tensor x) {
+  return Rice::detail::To_Ruby<torch::Tensor>().convert(x);
 }
 
-inline Object wrap(torch::Scalar x) {
-  return to_ruby<torch::Scalar>(x);
+inline VALUE wrap(torch::Scalar x) {
+  return Rice::detail::To_Ruby<torch::Scalar>().convert(x);
 }
 
-inline Object wrap(torch::ScalarType x) {
-  return to_ruby<torch::ScalarType>(x);
+inline VALUE wrap(torch::ScalarType x) {
+  return Rice::detail::To_Ruby<torch::ScalarType>().convert(x);
 }
 
-inline Object wrap(torch::QScheme x) {
-  return to_ruby<torch::QScheme>(x);
+inline VALUE wrap(torch::QScheme x) {
+  return Rice::detail::To_Ruby<torch::QScheme>().convert(x);
 }
 
-inline Object wrap(std::tuple<torch::Tensor, torch::Tensor> x) {
-  Array a;
-  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
-  return Object(a);
+inline VALUE wrap(std::tuple<torch::Tensor, torch::Tensor> x) {
+  return rb_ary_new3(
+    2,
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<0>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<1>(x))
+  );
 }
 
-inline Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> x) {
-  Array a;
-  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<2>(x)));
-  return Object(a);
+inline VALUE wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> x) {
+  return rb_ary_new3(
+    3,
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<0>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<1>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<2>(x))
+  );
 }
 
-inline Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> x) {
-  Array a;
-  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<2>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<3>(x)));
-  return Object(a);
+inline VALUE wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> x) {
+  return rb_ary_new3(
+    4,
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<0>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<1>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<2>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<3>(x))
+  );
 }
 
-inline Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> x) {
-  Array a;
-  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<2>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<3>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<4>(x)));
-  return Object(a);
+inline VALUE wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> x) {
+  return rb_ary_new3(
+    5,
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<0>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<1>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<2>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<3>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<4>(x))
+  );
 }
 
-inline Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, int64_t> x) {
-  Array a;
-  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<2>(x)));
-  a.push(to_ruby<int64_t>(std::get<3>(x)));
-  return Object(a);
+inline VALUE wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, int64_t> x) {
+  return rb_ary_new3(
+    4,
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<0>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<1>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<2>(x)),
+    Rice::detail::To_Ruby<int64_t>().convert(std::get<3>(x))
+  );
 }
 
-inline Object wrap(std::tuple<torch::Tensor, torch::Tensor, double, int64_t> x) {
-  Array a;
-  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
-  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
-  a.push(to_ruby<double>(std::get<2>(x)));
-  a.push(to_ruby<int64_t>(std::get<3>(x)));
-  return Object(a);
+inline VALUE wrap(std::tuple<torch::Tensor, torch::Tensor, double, int64_t> x) {
+  return rb_ary_new3(
+    4,
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<0>(x)),
+    Rice::detail::To_Ruby<torch::Tensor>().convert(std::get<1>(x)),
+    Rice::detail::To_Ruby<double>().convert(std::get<2>(x)),
+    Rice::detail::To_Ruby<int64_t>().convert(std::get<3>(x))
+  );
 }
 
-inline Object wrap(torch::TensorList x) {
-  Array a;
-  for (auto& t : x) {
-    a.push(to_ruby<torch::Tensor>(t));
+inline VALUE wrap(torch::TensorList x) {
+  auto a = rb_ary_new2(x.size());
+  for (auto t : x) {
+    rb_ary_push(a, Rice::detail::To_Ruby<torch::Tensor>().convert(t));
   }
-  return Object(a);
+  return a;
 }
 
-inline Object wrap(std::tuple<double, double> x) {
-  Array a;
-  a.push(to_ruby<double>(std::get<0>(x)));
-  a.push(to_ruby<double>(std::get<1>(x)));
-  return Object(a);
+inline VALUE wrap(std::tuple<double, double> x) {
+  return rb_ary_new3(
+    2,
+    Rice::detail::To_Ruby<double>().convert(std::get<0>(x)),
+    Rice::detail::To_Ruby<double>().convert(std::get<1>(x))
+  );
 }

data/lib/torch.rb

@@ -238,8 +238,11 @@ module Torch
     double: 7,
     float64: 7,
     complex_half: 8,
+    complex32: 8,
     complex_float: 9,
+    complex64: 9,
     complex_double: 10,
+    complex128: 10,
     bool: 11,
     qint8: 12,
     quint8: 13,
@@ -337,25 +340,24 @@ module Torch
       }
     end
 
-    def no_grad
-      previous_value = grad_enabled?
-      begin
-        _set_grad_enabled(false)
-        yield
-      ensure
-        _set_grad_enabled(previous_value)
-      end
+    def no_grad(&block)
+      grad_enabled(false, &block)
+    end
+
+    def enable_grad(&block)
+      grad_enabled(true, &block)
     end
 
-    def enable_grad
+    def grad_enabled(value)
       previous_value = grad_enabled?
       begin
-        _set_grad_enabled(true)
+        _set_grad_enabled(value)
         yield
       ensure
         _set_grad_enabled(previous_value)
       end
     end
+    alias_method :set_grad_enabled, :grad_enabled
 
     def device(str)
       Device.new(str)
@@ -395,6 +397,8 @@ module Torch
           options[:dtype] = :int64
         elsif data.all? { |v| v == true || v == false }
           options[:dtype] = :bool
+        elsif data.any? { |v| v.is_a?(Complex) }
+          options[:dtype] = :complex64
         end
       end
 

data/lib/torch/inspector.rb

@@ -96,8 +96,11 @@ module Torch
             ret = "%.#{PRINT_OPTS[:precision]}f" % value
           end
         elsif @complex_dtype
-          p = PRINT_OPTS[:precision]
-          raise NotImplementedYet
+          # TODO use float formatter for each part
+          precision = PRINT_OPTS[:precision]
+          imag = value.imag
+          sign = imag >= 0 ? "+" : "-"
+          ret = "%.#{precision}f#{sign}%.#{precision}fi" % [value.real, value.imag.abs]
         else
           ret = value.to_s
         end

data/lib/torch/nn/linear.rb

@@ -1,6 +1,8 @@
 module Torch
   module NN
     class Linear < Module
+      attr_reader :in_features, :out_features
+
       def initialize(in_features, out_features, bias: true)
         super()
         @in_features = in_features

data/lib/torch/nn/module.rb

@@ -113,35 +113,53 @@ module Torch
         forward(*input, **kwargs)
       end
 
-      def state_dict(destination: nil)
+      def state_dict(destination: nil, prefix: "")
         destination ||= {}
-        named_parameters.each do |k, v|
-          destination[k] = v
+        save_to_state_dict(destination, prefix: prefix)
+
+        named_children.each do |name, mod|
+          next unless mod
+          mod.state_dict(destination: destination, prefix: prefix + name + ".")
         end
         destination
       end
 
-      # TODO add strict option
-      # TODO match PyTorch behavior
-      def load_state_dict(state_dict)
-        state_dict.each do |k, input_param|
-          k1, k2 = k.split(".", 2)
-          mod = named_modules[k1]
-          if mod.is_a?(Module)
-            param = mod.named_parameters[k2]
-            if param.is_a?(Parameter)
-              Torch.no_grad do
-                param.copy!(input_param)
-              end
-            else
-              raise Error, "Unknown parameter: #{k1}"
-            end
-          else
-            raise Error, "Unknown module: #{k1}"
+      def load_state_dict(state_dict, strict: true)
+        # TODO support strict: false
+        raise "strict: false not implemented yet" unless strict
+
+        missing_keys = []
+        unexpected_keys = []
+        error_msgs = []
+
+        # TODO handle metadata
+
+        _load = lambda do |mod, prefix = ""|
+          # TODO handle metadata
+          local_metadata = {}
+          mod.send(:load_from_state_dict, state_dict, prefix, local_metadata, true, missing_keys, unexpected_keys, error_msgs)
+          mod.named_children.each do |name, child|
+            _load.call(child, prefix + name + ".") unless child.nil?
+          end
+        end
+
+        _load.call(self)
+
+        if strict
+          if unexpected_keys.any?
+            error_msgs << "Unexpected key(s) in state_dict: #{unexpected_keys.join(", ")}"
+          end
+
+          if missing_keys.any?
+            error_msgs << "Missing key(s) in state_dict: #{missing_keys.join(", ")}"
           end
         end
 
-        # TODO return missing keys and unexpected keys
+        if error_msgs.any?
+          # just show first error
+          raise Error, error_msgs[0]
+        end
+
         nil
       end
 
@@ -268,6 +286,12 @@ module Torch
           named_buffers[name]
         elsif named_modules.key?(name)
           named_modules[name]
+        elsif method.end_with?("=") && named_modules.key?(method[0..-2])
+          if instance_variable_defined?("@#{method[0..-2]}")
+            instance_variable_set("@#{method[0..-2]}", *args)
+          else
+            raise NotImplementedYet
+          end
         else
           super
         end
@@ -300,6 +324,68 @@ module Torch
       def dict
         instance_variables.reject { |k| instance_variable_get(k).is_a?(Tensor) }.map { |k| [k[1..-1].to_sym, instance_variable_get(k)] }.to_h
       end
+
+      def load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+        # TODO add hooks
+
+        # TODO handle non-persistent buffers
+        persistent_buffers = named_buffers
+        local_name_params = named_parameters(recurse: false).merge(persistent_buffers)
+        local_state = local_name_params.select { |_, v| !v.nil? }
+
+        local_state.each do |name, param|
+          key = prefix + name
+          if state_dict.key?(key)
+            input_param = state_dict[key]
+
+            # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+
+            if param.shape.length == 0 && input_param.shape.length == 1
+              input_param = input_param[0]
+            end
+
+            if input_param.shape != param.shape
+              # local shape should match the one in checkpoint
+              error_msgs << "size mismatch for #{key}: copying a param with shape #{input_param.shape} from checkpoint, " +
+                            "the shape in current model is #{param.shape}."
+              next
+            end
+
+            begin
+              Torch.no_grad do
+                param.copy!(input_param)
+              end
+            rescue => e
+              error_msgs << "While copying the parameter named #{key.inspect}, " +
+                            "whose dimensions in the model are #{param.size} and " +
+                            "whose dimensions in the checkpoint are #{input_param.size}, " +
+                            "an exception occurred: #{e.inspect}"
+            end
+          elsif strict
+            missing_keys << key
+          end
+        end
+
+        if strict
+          state_dict.each_key do |key|
+            if key.start_with?(prefix)
+              input_name = key[prefix.length..-1]
+              input_name = input_name.split(".", 2)[0]
+              if !named_children.key?(input_name) && !local_state.key?(input_name)
+                unexpected_keys << key
+              end
+            end
+          end
+        end
+      end
+
+      def save_to_state_dict(destination, prefix: "")
+        named_parameters(recurse: false).each do |k, v|
+          destination[prefix + k] = v
+        end
+        named_buffers.each do |k, v|
+          destination[prefix + k] = v
+        end
+      end
     end
   end
 end