torch-rb 0.8.2 → 0.9.2

data/ext/torch/backends.cpp

@@ -7,11 +7,11 @@
 void init_backends(Rice::Module& m) {
   auto rb_mBackends = Rice::define_module_under(m, "Backends");
 
-   Rice::define_module_under(rb_mBackends, "OpenMP")
+  Rice::define_module_under(rb_mBackends, "OpenMP")
     .add_handler<torch::Error>(handle_error)
     .define_singleton_function("available?", &torch::hasOpenMP);
 
-   Rice::define_module_under(rb_mBackends, "MKL")
+  Rice::define_module_under(rb_mBackends, "MKL")
     .add_handler<torch::Error>(handle_error)
     .define_singleton_function("available?", &torch::hasMKL);
 }

data/ext/torch/nn.cpp

@@ -98,8 +98,11 @@ void init_nn(Rice::Module& m) {
         auto grad = self.grad();
         return grad.defined() ? Object(Rice::detail::To_Ruby<torch::Tensor>().convert(grad)) : Nil;
       })
+    // can't use grad=
+    // assignment methods fail with Ruby 3.0
+    // TODO add checks like Tensor
     .define_method(
-      "grad=",
+      "_set_grad",
       [](Parameter& self, torch::Tensor& grad) {
         self.mutable_grad() = grad;
       })

data/ext/torch/ruby_arg_parser.cpp

@@ -472,12 +472,12 @@ static void extra_kwargs(FunctionSignature& signature, VALUE kwargs, ssize_t num
     auto param_idx = find_param(signature, key);
     if (param_idx < 0) {
       rb_raise(rb_eArgError, "%s() got an unexpected keyword argument '%s'",
-          signature.name.c_str(), THPUtils_unpackSymbol(key).c_str());
+          signature.name.c_str(), rb_id2name(rb_to_id(key)));
     }
 
     if (param_idx < num_pos_args) {
       rb_raise(rb_eArgError, "%s() got multiple values for argument '%s'",
-          signature.name.c_str(), THPUtils_unpackSymbol(key).c_str());
+          signature.name.c_str(), rb_id2name(rb_to_id(key)));
     }
   }
 

data/ext/torch/ruby_arg_parser.h

@@ -75,7 +75,7 @@ struct RubyArgs {
   int idx;
 
   inline at::Tensor tensor(int i);
-  inline OptionalTensor optionalTensor(int i);
+  inline c10::optional<at::Tensor> optionalTensor(int i);
   inline at::Scalar scalar(int i);
   // inline at::Scalar scalarWithDefault(int i, at::Scalar default_scalar);
   inline std::vector<at::Scalar> scalarlist(int i);
@@ -109,6 +109,9 @@ struct RubyArgs {
   // inline at::QScheme toQScheme(int i);
   inline std::string string(int i);
   inline c10::optional<std::string> stringOptional(int i);
+  inline c10::string_view stringView(int i);
+  // inline c10::string_view stringViewWithDefault(int i, const c10::string_view default_str);
+  inline c10::optional<c10::string_view> stringViewOptional(int i);
   // inline PyObject* pyobject(int i);
   inline int64_t toInt64(int i);
   // inline int64_t toInt64WithDefault(int i, int64_t default_int);
@@ -125,8 +128,8 @@ inline at::Tensor RubyArgs::tensor(int i) {
   return Rice::detail::From_Ruby<torch::Tensor>().convert(args[i]);
 }
 
-inline OptionalTensor RubyArgs::optionalTensor(int i) {
-  if (NIL_P(args[i])) return OptionalTensor(Nil);
+inline c10::optional<at::Tensor> RubyArgs::optionalTensor(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
   return tensor(i);
 }
 
@@ -232,7 +235,7 @@ inline ScalarType RubyArgs::scalartype(int i) {
 
   auto it = dtype_map.find(args[i]);
   if (it == dtype_map.end()) {
-    rb_raise(rb_eArgError, "invalid dtype: %s", THPUtils_unpackSymbol(args[i]).c_str());
+    rb_raise(rb_eArgError, "invalid dtype: %s", rb_id2name(rb_to_id(args[i])));
   }
   return it->second;
 }
@@ -290,7 +293,7 @@ inline c10::optional<at::Layout> RubyArgs::layoutOptional(int i) {
 
   auto it = layout_map.find(args[i]);
   if (it == layout_map.end()) {
-    rb_raise(rb_eArgError, "invalid layout: %s", THPUtils_unpackSymbol(args[i]).c_str());
+    rb_raise(rb_eArgError, "invalid layout: %s", rb_id2name(rb_to_id(args[i])));
   }
   return it->second;
 }
@@ -322,6 +325,17 @@ inline c10::optional<std::string> RubyArgs::stringOptional(int i) {
   return Rice::detail::From_Ruby<std::string>().convert(args[i]);
 }
 
+// string_view does not own data
+inline c10::string_view RubyArgs::stringView(int i) {
+  return c10::string_view(RSTRING_PTR(args[i]), RSTRING_LEN(args[i]));
+}
+
+// string_view does not own data
+inline c10::optional<c10::string_view> RubyArgs::stringViewOptional(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
+  return c10::string_view(RSTRING_PTR(args[i]), RSTRING_LEN(args[i]));
+}
+
 inline int64_t RubyArgs::toInt64(int i) {
   if (NIL_P(args[i])) return signature.params[i].default_int;
   return Rice::detail::From_Ruby<int64_t>().convert(args[i]);

data/ext/torch/templates.h

@@ -41,23 +41,39 @@ using torch::nn::init::NonlinearityType;
 #define RETURN_NIL                                                   \
   return Qnil;
 
-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());
-      }
+namespace Rice::detail
+{
+  template<typename T>
+  struct Type<c10::complex<T>>
+  {
+    static bool verify()
+    {
+      return true;
     }
-    OptionalTensor(torch::Tensor o) {
-      value = o;
+  };
+
+  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());
     }
-    operator torch::Tensor() const {
-      return value;
+  };
+
+  template<typename T>
+  class From_Ruby<c10::complex<T>>
+  {
+  public:
+    c10::complex<T> convert(VALUE x)
+    {
+      VALUE real = rb_funcall(x, rb_intern("real"), 0);
+      VALUE imag = rb_funcall(x, rb_intern("imag"), 0);
+      return c10::complex<T>(From_Ruby<T>().convert(real), From_Ruby<T>().convert(imag));
     }
-};
+  };
+}
 
 namespace Rice::detail
 {
@@ -131,25 +147,6 @@ namespace Rice::detail
     }
   };
 
-  template<>
-  struct Type<OptionalTensor>
-  {
-    static bool verify()
-    {
-      return true;
-    }
-  };
-
-  template<>
-  class From_Ruby<OptionalTensor>
-  {
-  public:
-    OptionalTensor convert(VALUE x)
-    {
-      return OptionalTensor(x);
-    }
-  };
-
   template<>
   struct Type<Scalar>
   {

data/ext/torch/tensor.cpp

@@ -10,28 +10,6 @@
 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()});
@@ -107,7 +85,7 @@ static VALUE tensor__backward(int argc, VALUE* argv, VALUE self_)
   ParsedArgs<4> parsed_args;
   auto _r = parser.parse(self_, argc, argv, parsed_args);
   // _backward(Tensor self, Tensor[] inputs, Tensor? gradient=None, bool? retain_graph=None, bool create_graph=False) -> ()
-  auto dispatch__backward = [](const Tensor & self, TensorList inputs, const OptionalTensor & gradient, c10::optional<bool> retain_graph, bool create_graph) -> void {
+  auto dispatch__backward = [](const Tensor & self, TensorList inputs, const c10::optional<at::Tensor> & gradient, c10::optional<bool> retain_graph, bool create_graph) -> void {
     // in future, release GVL
     self._backward(inputs, gradient, retain_graph, create_graph);
   };
@@ -125,13 +103,13 @@ void init_tensor(Rice::Module& m, Rice::Class& c, Rice::Class& rb_cTensorOptions
   rb_define_method(rb_cTensor, "backward", (VALUE (*)(...)) tensor__backward, -1);
 
   rb_cTensor
-    .define_method("cuda?", &torch::Tensor::is_cuda)
-    .define_method("sparse?", &torch::Tensor::is_sparse)
-    .define_method("quantized?", &torch::Tensor::is_quantized)
-    .define_method("dim", &torch::Tensor::dim)
-    .define_method("numel", &torch::Tensor::numel)
-    .define_method("element_size", &torch::Tensor::element_size)
-    .define_method("requires_grad", &torch::Tensor::requires_grad)
+    .define_method("cuda?", [](Tensor& self) { return self.is_cuda(); })
+    .define_method("sparse?", [](Tensor& self) { return self.is_sparse(); })
+    .define_method("quantized?", [](Tensor& self) { return self.is_quantized(); })
+    .define_method("dim", [](Tensor& self) { return self.dim(); })
+    .define_method("numel", [](Tensor& self) { return self.numel(); })
+    .define_method("element_size", [](Tensor& self) { return self.element_size(); })
+    .define_method("requires_grad", [](Tensor& self) { return self.requires_grad(); })
     .define_method(
       "_size",
       [](Tensor& self, int64_t dim) {
@@ -189,9 +167,31 @@ void init_tensor(Rice::Module& m, Rice::Class& c, Rice::Class& rb_cTensorOptions
         auto grad = self.grad();
         return grad.defined() ? Object(Rice::detail::To_Ruby<torch::Tensor>().convert(grad)) : Nil;
       })
+    // can't use grad=
+    // assignment methods fail with Ruby 3.0
     .define_method(
-      "grad=",
-      [](Tensor& self, torch::Tensor& grad) {
+      "_set_grad",
+      [](Tensor& self, Rice::Object value) {
+        if (value.is_nil()) {
+          self.mutable_grad().reset();
+          return;
+        }
+
+        const auto& grad = Rice::detail::From_Ruby<torch::Tensor>().convert(value.value());
+
+        // TODO support sparse grad
+        if (!grad.options().type_equal(self.options())) {
+          rb_raise(rb_eArgError, "assigned grad has data of a different type");
+        }
+
+        if (self.is_cuda() && grad.get_device() != self.get_device()) {
+          rb_raise(rb_eArgError, "assigned grad has data located on a different device");
+        }
+
+        if (!self.sizes().equals(grad.sizes())) {
+          rb_raise(rb_eArgError, "assigned grad has data of a different size");
+        }
+
         self.mutable_grad() = grad;
       })
     .define_method(
@@ -281,7 +281,7 @@ void init_tensor(Rice::Module& m, Rice::Class& c, Rice::Class& rb_cTensorOptions
       })
     .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);
       });
 

data/ext/torch/torch.cpp

@@ -6,6 +6,23 @@
 #include "templates.h"
 #include "utils.h"
 
+template<typename T>
+torch::Tensor make_tensor(Rice::Array a, std::vector<int64_t> size, const torch::TensorOptions &options) {
+  std::vector<T> vec;
+  for (long i = 0; i < a.size(); i++) {
+    vec.push_back(Rice::detail::From_Ruby<T>().convert(a[i].value()));
+  }
+
+  // hack for requires_grad error
+  auto requires_grad = options.requires_grad();
+  torch::Tensor t = torch::tensor(vec, options.requires_grad(c10::nullopt));
+  if (requires_grad) {
+    t.set_requires_grad(true);
+  }
+
+  return t.reshape(size);
+}
+
 void init_torch(Rice::Module& m) {
   m.add_handler<torch::Error>(handle_error);
   add_torch_functions(m);
@@ -61,35 +78,28 @@ void init_torch(Rice::Module& m) {
       "_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);
+        if (dtype == torch::kByte) {
+          return make_tensor<uint8_t>(a, size, options);
+        } else if (dtype == torch::kChar) {
+          return make_tensor<int8_t>(a, size, options);
+        } else if (dtype == torch::kShort) {
+          return make_tensor<int16_t>(a, size, options);
+        } else if (dtype == torch::kInt) {
+          return make_tensor<int32_t>(a, size, options);
+        } else if (dtype == torch::kLong) {
+          return make_tensor<int64_t>(a, size, options);
+        } else if (dtype == torch::kFloat) {
+          return make_tensor<float>(a, size, options);
+        } else if (dtype == torch::kDouble) {
+          return make_tensor<double>(a, size, options);
+        } else if (dtype == torch::kBool) {
+          return make_tensor<uint8_t>(a, size, options);
+        } else if (dtype == torch::kComplexFloat) {
+          return make_tensor<c10::complex<float>>(a, size, options);
+        } else if (dtype == torch::kComplexDouble) {
+          return make_tensor<c10::complex<double>>(a, size, 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);
-          }
+          throw std::runtime_error("Unsupported type");
         }
-        return t.reshape(size);
       });
 }

data/ext/torch/utils.h

@@ -16,12 +16,6 @@ inline VALUE THPUtils_internSymbol(const std::string& str) {
   return Rice::Symbol(str);
 }
 
-inline std::string THPUtils_unpackSymbol(VALUE obj) {
-  Check_Type(obj, T_SYMBOL);
-  obj = rb_funcall(obj, rb_intern("to_s"), 0);
-  return std::string(RSTRING_PTR(obj), RSTRING_LEN(obj));
-}
-
 inline std::string THPUtils_unpackString(VALUE obj) {
   Check_Type(obj, T_STRING);
   return std::string(RSTRING_PTR(obj), RSTRING_LEN(obj));

data/lib/torch/inspector.rb

@@ -247,7 +247,7 @@ module Torch
       # length includes spaces and comma between elements
       element_length = formatter.width + 2
       elements_per_line = [1, ((PRINT_OPTS[:linewidth] - indent) / element_length.to_f).floor.to_i].max
-      char_per_line = element_length * elements_per_line
+      _char_per_line = element_length * elements_per_line
 
       if summarize && slf.size(0) > 2 * PRINT_OPTS[:edgeitems]
         data = (

data/lib/torch/nn/functional.rb

@@ -571,7 +571,7 @@ module Torch
         end
 
         def _interp_output_size(closed_over_args)
-          input, size, scale_factor, recompute_scale_factor = closed_over_args
+          input, size, scale_factor, _recompute_scale_factor = closed_over_args
           dim = input.dim - 2
           if size.nil? && scale_factor.nil?
             raise ArgumentError, "either size or scale_factor should be defined"

data/lib/torch/nn/functional_attention.rb

@@ -56,7 +56,7 @@ module Torch
           attn_mask: nil, dropout_p: 0.0
         )
 
-          b, nt, e = q.shape
+          _b, _nt, e = q.shape
 
           q = q / Math.sqrt(e)
 

data/lib/torch/nn/module.rb

@@ -280,6 +280,11 @@ module Torch
         end
       end
 
+      def deep_dup
+        memo = {}
+        dup_value(self, memo)
+      end
+
       def method_missing(method, *args, &block)
         name = method.to_s
         if named_parameters.key?(name)
@@ -388,6 +393,29 @@ module Torch
           destination[prefix + k] = v
         end
       end
+
+      # keep memo hash like Python deepcopy
+      # https://docs.python.org/3/library/copy.html
+      def dup_value(v, memo)
+        memo[v.object_id] ||= begin
+          case v
+          when Method, UnboundMethod
+            v
+          when Hash
+            v.to_h { |k, v2| [dup_value(k, memo), dup_value(v2, memo)] }
+          when Array
+            v.map { |v2| dup_value(v2, memo) }
+          when Torch::NN::Module
+            copy = v.dup
+            v.instance_variables.each do |var|
+              copy.instance_variable_set(var, dup_value(v.instance_variable_get(var), memo))
+            end
+            copy
+          else
+            v.dup
+          end
+        end
+      end
     end
   end
 end

data/lib/torch/nn/parameter.rb

@@ -1,6 +1,9 @@
 module Torch
   module NN
     class Parameter < Tensor
+      # fix for issue w/ assignment methods
+      alias_method :grad=, :_set_grad
+
       def self.new(data = nil, requires_grad: true)
         data = Tensor.new unless data
         _make_subclass(data, requires_grad)
@@ -9,6 +12,12 @@ module Torch
       def inspect
         "Parameter containing:\n#{super}"
       end
+
+      def dup
+        Torch.no_grad do
+          Parameter.new(clone, requires_grad: requires_grad)
+        end
+      end
     end
   end
 end

data/lib/torch/nn/transformer_decoder_layer.rb

@@ -35,7 +35,7 @@ module Torch
         tgt += @dropout2.(tgt2)
         tgt = @norm2.(tgt)
         tgt2 = @linear2.(@dropout.(@activation.(@linear1.(tgt))))
-        tgt += @dropout3.(tgt)
+        tgt += @dropout3.(tgt2)
         @norm3.(tgt)
       end
     end

data/lib/torch/nn/utils.rb

@@ -22,11 +22,7 @@ module Torch
       end
 
       def _clones(mod, n)
-        state = mod.state_dict
-        layers = n.times.map do |i|
-          mod.clone.tap { |l| l.load_state_dict(state) }
-        end
-        ModuleList.new(layers)
+        ModuleList.new(n.times.map { mod.deep_dup })
       end
 
       def _activation_fn(activation)

data/lib/torch/tensor.rb

@@ -8,6 +8,9 @@ module Torch
     alias_method :ndim, :dim
     alias_method :ndimension, :dim
 
+    # fix for issue w/ assignment methods
+    alias_method :grad=, :_set_grad
+
     # use alias_method for performance
     alias_method :+, :add
     alias_method :-, :sub
@@ -106,6 +109,7 @@ module Torch
       size(0)
     end
 
+    remove_method :item
     def item
       if numel != 1
         raise Error, "only one element tensors can be converted to Ruby scalars"
@@ -133,18 +137,10 @@ module Torch
       cls.from_string(_data_str).reshape(*shape)
     end
 
-    def new_ones(*size, **options)
-      Torch.ones_like(Torch.empty(*size), **options)
-    end
-
     def requires_grad=(requires_grad)
       _requires_grad!(requires_grad)
     end
 
-    def requires_grad!(requires_grad = true)
-      _requires_grad!(requires_grad)
-    end
-
     def type(dtype)
       if dtype.is_a?(Class)
         raise Error, "Invalid type: #{dtype}" unless TENSOR_TYPE_CLASSES.include?(dtype)
@@ -185,5 +181,23 @@ module Torch
     def stft(*args)
       Torch.stft(*args)
     end
+
+    def dup
+      Torch.no_grad do
+        clone
+      end
+    end
+
+    # not a method in native_functions.yaml
+    # attribute in Python rather than method
+    def imag
+      Torch.imag(self)
+    end
+
+    # not a method in native_functions.yaml
+    # attribute in Python rather than method
+    def real
+      Torch.real(self)
+    end
   end
 end

data/lib/torch/utils/data/data_loader.rb

@@ -29,7 +29,7 @@ module Torch
 
           # try to keep the random number generator in sync with Python
           # this makes it easy to compare results
-          base_seed = Torch.empty([], dtype: :int64).random!.item
+          _base_seed = Torch.empty([], dtype: :int64).random!.item
 
           indexes =
             if @shuffle

data/lib/torch/version.rb

@@ -1,3 +1,3 @@
 module Torch
-  VERSION = "0.8.2"
+  VERSION = "0.9.2"
 end