torch-rb 0.20.0 → 0.21.0

data/ext/torch/device.cpp

@@ -1,6 +1,9 @@
+#include <string>
+
 #include <torch/torch.h>
 
 #include <rice/rice.hpp>
+#include <rice/stl.hpp>
 
 #include "utils.h"
 

data/ext/torch/ext.cpp

@@ -17,8 +17,7 @@ void init_ivalue(Rice::Module& m, Rice::Class& rb_cIValue);
 void init_random(Rice::Module& m);
 
 extern "C"
-void Init_ext()
-{
+void Init_ext() {
   auto m = Rice::define_module("Torch");
 
   // need to define certain classes up front to keep Rice happy

data/ext/torch/ivalue.cpp

@@ -1,3 +1,5 @@
+#include <utility>
+
 #include <torch/torch.h>
 
 #include <rice/rice.hpp>

data/ext/torch/nn.cpp

@@ -1,3 +1,5 @@
+#include <utility>
+
 #include <torch/torch.h>
 
 #include <rice/rice.hpp>
@@ -93,7 +95,7 @@ void init_nn(Rice::Module& m) {
       "grad",
       [](Parameter& self) {
         auto grad = self.grad();
-        return grad.defined() ? Object(Rice::detail::To_Ruby<torch::Tensor>().convert(grad)) : Nil;
+        return grad.defined() ? Rice::Object(Rice::detail::To_Ruby<torch::Tensor>().convert(grad)) : Rice::Nil;
       })
     // can't use grad=
     // assignment methods fail with Ruby 3.0

data/ext/torch/ruby_arg_parser.cpp

@@ -1,5 +1,10 @@
 // adapted from PyTorch - python_arg_parser.cpp
 
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
 #include "ruby_arg_parser.h"
 
 VALUE THPGeneratorClass = Qnil;
@@ -99,7 +104,7 @@ FunctionParameter::FunctionParameter(const std::string& fmt, bool keyword_only)
   ruby_name = THPUtils_internSymbol(name);
   auto np_compat_it = numpy_compatibility_arg_names.find(name);
   if (np_compat_it != numpy_compatibility_arg_names.end()) {
-    for (const auto& str: np_compat_it->second) {
+    for (const auto& str : np_compat_it->second) {
       numpy_python_names.push_back(THPUtils_internSymbol(str));
     }
   }
@@ -190,8 +195,7 @@ static bool is_int_or_symint_list(VALUE obj, int broadcast_size) {
 }
 
 // argnum is needed for raising the TypeError, it's used in the error message.
-auto FunctionParameter::check(VALUE obj, int argnum) -> bool
-{
+auto FunctionParameter::check(VALUE obj, int argnum) -> bool {
   switch (type_) {
     case ParameterType::TENSOR: {
       if (THPVariable_Check(obj)) {

data/ext/torch/ruby_arg_parser.h

@@ -3,6 +3,9 @@
 #pragma once
 
 #include <sstream>
+#include <unordered_map>
+#include <string>
+#include <vector>
 
 #include <torch/torch.h>
 #include <rice/rice.hpp>
@@ -162,7 +165,7 @@ inline std::array<at::Tensor, N> RubyArgs::tensorlist_n(int i) {
   Check_Type(arg, T_ARRAY);
   auto size = RARRAY_LEN(arg);
   if (size != N) {
-    rb_raise(rb_eArgError, "expected array of %d elements but got %d", N, (int)size);
+    rb_raise(rb_eArgError, "expected array of %d elements but got %d", N, static_cast<int>(size));
   }
   for (int idx = 0; idx < size; idx++) {
     VALUE obj = rb_ary_entry(arg, idx);
@@ -463,7 +466,7 @@ struct RubyArgParser {
     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);
+        rb_raise(rb_eArgError, "RubyArgParser: dst ParsedArgs buffer does not have enough capacity, expected %d (got %d)", static_cast<int>(max_args), N);
       }
       return raw_parse(self, argc, argv, dst.args);
     }

data/ext/torch/templates.h

@@ -1,13 +1,13 @@
 #pragma once
 
+#include <string>
+
 #ifdef isfinite
 #undef isfinite
 #endif
 
 #include <rice/rice.hpp>
 
-using namespace Rice;
-
 using torch::Device;
 using torch::Scalar;
 using torch::ScalarType;
@@ -41,55 +41,43 @@ using torch::nn::init::NonlinearityType;
 #define RETURN_NIL                                                   \
   return Qnil;
 
-namespace Rice::detail
-{
+namespace Rice::detail {
   template<typename T>
-  struct Type<c10::complex<T>>
-  {
+  struct Type<c10::complex<T>> {
     static bool verify() { return true; }
   };
 
   template<typename T>
-  class To_Ruby<c10::complex<T>>
-  {
+  class To_Ruby<c10::complex<T>> {
   public:
-    VALUE convert(c10::complex<T> const& x)
-    {
+    VALUE convert(c10::complex<T> const& x) {
       return rb_dbl_complex_new(x.real(), x.imag());
     }
   };
 
   template<typename T>
-  class From_Ruby<c10::complex<T>>
-  {
+  class From_Ruby<c10::complex<T>> {
   public:
     Convertible is_convertible(VALUE value) { return Convertible::Cast; }
 
-    c10::complex<T> convert(VALUE x)
-    {
+    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
-{
   template<>
-  struct Type<FanModeType>
-  {
+  struct Type<FanModeType> {
     static bool verify() { return true; }
   };
 
   template<>
-  class From_Ruby<FanModeType>
-  {
+  class From_Ruby<FanModeType> {
   public:
     Convertible is_convertible(VALUE value) { return Convertible::Cast; }
 
-    FanModeType convert(VALUE x)
-    {
+    FanModeType convert(VALUE x) {
       auto s = String(x).str();
       if (s == "fan_in") {
         return torch::kFanIn;
@@ -102,19 +90,16 @@ namespace Rice::detail
   };
 
   template<>
-  struct Type<NonlinearityType>
-  {
+  struct Type<NonlinearityType> {
     static bool verify() { return true; }
   };
 
   template<>
-  class From_Ruby<NonlinearityType>
-  {
+  class From_Ruby<NonlinearityType> {
   public:
     Convertible is_convertible(VALUE value) { return Convertible::Cast; }
 
-    NonlinearityType convert(VALUE x)
-    {
+    NonlinearityType convert(VALUE x) {
       auto s = String(x).str();
       if (s == "linear") {
         return torch::kLinear;
@@ -145,19 +130,16 @@ namespace Rice::detail
   };
 
   template<>
-  struct Type<Scalar>
-  {
+  struct Type<Scalar> {
     static bool verify() { return true; }
   };
 
   template<>
-  class From_Ruby<Scalar>
-  {
+  class From_Ruby<Scalar> {
   public:
     Convertible is_convertible(VALUE value) { return Convertible::Cast; }
 
-    Scalar convert(VALUE x)
-    {
+    Scalar convert(VALUE x) {
       if (FIXNUM_P(x)) {
         return torch::Scalar(From_Ruby<int64_t>().convert(x));
       } else {
@@ -165,4 +147,4 @@ namespace Rice::detail
       }
     }
   };
-}
+} // namespace Rice::detail

data/ext/torch/tensor.cpp

@@ -1,3 +1,6 @@
+#include <string>
+#include <vector>
+
 #include <torch/torch.h>
 
 #include <rice/rice.hpp>
@@ -7,7 +10,8 @@
 #include "templates.h"
 #include "utils.h"
 
-using namespace Rice;
+using Rice::Array;
+using Rice::Object;
 using torch::indexing::TensorIndex;
 
 template<typename T>
@@ -21,7 +25,7 @@ Array flat_data(Tensor& tensor) {
   return a;
 }
 
-Class rb_cTensor;
+Rice::Class rb_cTensor;
 
 std::vector<TensorIndex> index_vector(Array a) {
   Object obj;
@@ -62,10 +66,10 @@ std::vector<TensorIndex> index_vector(Array a) {
       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) {
+    } else if (obj == Rice::True || obj == Rice::False) {
       indices.push_back(Rice::detail::From_Ruby<bool>().convert(obj.value()));
     } else {
-      throw Exception(rb_eArgError, "Unsupported index type: %s", rb_obj_classname(obj));
+      throw Rice::Exception(rb_eArgError, "Unsupported index type: %s", rb_obj_classname(obj));
     }
   }
   return indices;
@@ -75,8 +79,7 @@ std::vector<TensorIndex> index_vector(Array a) {
 // https://github.com/pytorch/pytorch/commit/2e5bfa9824f549be69a28e4705a72b4cf8a4c519
 // TODO add support for inputs argument
 // _backward
-static VALUE tensor__backward(int argc, VALUE* argv, VALUE self_)
-{
+static VALUE tensor__backward(int argc, VALUE* argv, VALUE self_) {
   HANDLE_TH_ERRORS
   Tensor& self = Rice::detail::From_Ruby<Tensor&>().convert(self_);
   static RubyArgParser parser({
@@ -165,7 +168,7 @@ void init_tensor(Rice::Module& m, Rice::Class& c, Rice::Class& rb_cTensorOptions
       "grad",
       [](Tensor& self) {
         auto grad = self.grad();
-        return grad.defined() ? Object(Rice::detail::To_Ruby<torch::Tensor>().convert(grad)) : Nil;
+        return grad.defined() ? Object(Rice::detail::To_Ruby<torch::Tensor>().convert(grad)) : Rice::Nil;
       })
     // can't use grad=
     // assignment methods fail with Ruby 3.0
@@ -197,7 +200,7 @@ void init_tensor(Rice::Module& m, Rice::Class& c, Rice::Class& rb_cTensorOptions
     .define_method(
       "_dtype",
       [](Tensor& self) {
-        return (int) at::typeMetaToScalarType(self.dtype());
+        return static_cast<int>(at::typeMetaToScalarType(self.dtype()));
       })
     .define_method(
       "_type",

data/ext/torch/torch.cpp

@@ -1,8 +1,11 @@
+#include <fstream>
+#include <string>
+#include <vector>
+
 #include <torch/torch.h>
 
 #include <rice/rice.hpp>
-
-#include <fstream>
+#include <rice/stl.hpp>
 
 #include "torch_functions.h"
 #include "templates.h"
@@ -60,7 +63,7 @@ void init_torch(Rice::Module& m) {
       "_save",
       [](const torch::IValue &value) {
         auto v = torch::pickle_save(value);
-        return Object(rb_str_new(v.data(), v.size()));
+        return Rice::Object(rb_str_new(v.data(), v.size()));
       })
     .define_singleton_function(
       "_load",

data/ext/torch/utils.h

@@ -1,12 +1,14 @@
 #pragma once
 
+#include <string>
+
 #include <torch/torch.h>
 
 #include <rice/rice.hpp>
 #include <rice/stl.hpp>
 
 static_assert(
-  TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 7,
+  TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 8,
   "Incompatible LibTorch version"
 );
 

data/lib/torch/nn/conv1d.rb

@@ -1,9 +1,17 @@
 module Torch
   module NN
     class Conv1d < ConvNd
-      def initialize(in_channels, out_channels, kernel_size, stride: 1,
-        padding: 0, dilation: 1, groups: 1, bias: true, padding_mode: "zeros")
-
+      def initialize(
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride: 1,
+        padding: 0,
+        dilation: 1,
+        groups: 1,
+        bias: true,
+        padding_mode: "zeros"
+      )
         kernel_size = _single(kernel_size)
         stride = _single(stride)
         padding = _single(padding)

data/lib/torch/nn/conv2d.rb

@@ -1,9 +1,17 @@
 module Torch
   module NN
     class Conv2d < ConvNd
-      def initialize(in_channels, out_channels, kernel_size, stride: 1,
-        padding: 0, dilation: 1, groups: 1, bias: true, padding_mode: "zeros")
-
+      def initialize(
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride: 1,
+        padding: 0,
+        dilation: 1,
+        groups: 1,
+        bias: true,
+        padding_mode: "zeros"
+      )
         kernel_size = _pair(kernel_size)
         stride = _pair(stride)
         padding = _pair(padding)

data/lib/torch/nn/conv3d.rb

@@ -1,9 +1,17 @@
 module Torch
   module NN
     class Conv3d < ConvNd
-      def initialize(in_channels, out_channels, kernel_size, stride: 1,
-        padding: 0, dilation: 1, groups: 1, bias: true, padding_mode: "zeros")
-
+      def initialize(
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride: 1,
+        padding: 0,
+        dilation: 1,
+        groups: 1,
+        bias: true,
+        padding_mode: "zeros"
+      )
         kernel_size = _triple(kernel_size)
         stride = _triple(stride)
         padding = _triple(padding)

data/lib/torch/nn/convnd.rb

@@ -1,7 +1,7 @@
 module Torch
   module NN
     class ConvNd < Module
-      attr_reader :in_channels, :out_channels, :kernel_size, :stride, :padding, :dilation, :transposed, :output_paddding, :groups, :padding_mode
+      attr_reader :in_channels, :out_channels, :kernel_size, :stride, :padding, :dilation, :transposed, :output_padding, :groups, :padding_mode
 
       def initialize(in_channels, out_channels, kernel_size, stride, padding, dilation, transposed, output_padding, groups, bias, padding_mode)
         super()

data/lib/torch/nn/embedding.rb

@@ -2,9 +2,16 @@
 module Torch
   module NN
     class Embedding < Module
-      def initialize(num_embeddings, embedding_dim, padding_idx: nil, max_norm: nil,
-        norm_type: 2.0, scale_grad_by_freq: false, sparse: false, _weight: nil)
-
+      def initialize(
+        num_embeddings,
+        embedding_dim,
+        padding_idx: nil,
+        max_norm: nil,
+        norm_type: 2.0,
+        scale_grad_by_freq: false,
+        sparse: false,
+        _weight: nil
+      )
         super()
         @num_embeddings = num_embeddings
         @embedding_dim = embedding_dim

data/lib/torch/nn/embedding_bag.rb

@@ -2,9 +2,16 @@
 module Torch
   module NN
     class EmbeddingBag < Module
-      def initialize(num_embeddings, embedding_dim, max_norm: nil, norm_type: 2.0,
-        scale_grad_by_freq: false, mode: "mean", sparse: false, _weight: nil)
-
+      def initialize(
+        num_embeddings,
+        embedding_dim,
+        max_norm: nil,
+        norm_type: 2.0,
+        scale_grad_by_freq: false,
+        mode: "mean",
+        sparse: false,
+        _weight: nil
+      )
         super()
         @num_embeddings = num_embeddings
         @embedding_dim = embedding_dim

data/lib/torch/nn/functional.rb

@@ -250,9 +250,16 @@ module Torch
 
         # normalization layers
 
-        def batch_norm(input, running_mean, running_var, weight: nil, bias: nil,
-          training: false, momentum: 0.1, eps: 1e-5)
-
+        def batch_norm(
+          input,
+          running_mean,
+          running_var,
+          weight: nil,
+          bias: nil,
+          training: false,
+          momentum: 0.1,
+          eps: 1e-5
+        )
           if training
             size = input.size
             size_prods = size[0]
@@ -274,9 +281,16 @@ module Torch
           Torch.group_norm(input, num_groups, weight, bias, eps, false)
         end
 
-        def instance_norm(input, running_mean: nil, running_var: nil, weight: nil,
-          bias: nil, use_input_stats: true, momentum: 0.1, eps: 1e-5)
-
+        def instance_norm(
+          input,
+          running_mean: nil,
+          running_var: nil,
+          weight: nil,
+          bias: nil,
+          use_input_stats: true,
+          momentum: 0.1,
+          eps: 1e-5
+        )
           Torch.instance_norm(
               input, weight, bias, running_mean, running_var,
               use_input_stats, momentum, eps, false

data/lib/torch/nn/functional_attention.rb

@@ -33,11 +33,16 @@ module Torch
         end
 
         def in_projection(
-          q, k, v,
-          w_q, w_k, w_v,
-          b_q: nil, b_k: nil, b_v: nil
+          q,
+          k,
+          v,
+          w_q,
+          w_k,
+          w_v,
+          b_q: nil,
+          b_k: nil,
+          b_v: nil
         )
-
           e_q, e_k, e_v = q.size(-1), k.size(-1), v.size(-1)
 
           raise ArgumentError, "Expecting query weights shape of #{[e_q, e_q]}, but got #{w_q.shape}" unless w_q.shape == [e_q, e_q]
@@ -52,10 +57,12 @@ module Torch
         end
 
         def scaled_dot_product_attention(
-          q, k, v,
-          attn_mask: nil, dropout_p: 0.0
+          q,
+          k,
+          v,
+          attn_mask: nil,
+          dropout_p: 0.0
         )
-
           _b, _nt, e = q.shape
 
           q = q / Math.sqrt(e)
@@ -71,22 +78,30 @@ module Torch
         end
 
         def multi_head_attention_forward(
-          query, key, value,
-          embed_dim_to_check, num_heads,
-          in_proj_weight, in_proj_bias,
-          bias_k, bias_v,
+          query,
+          key,
+          value,
+          embed_dim_to_check,
+          num_heads,
+          in_proj_weight,
+          in_proj_bias,
+          bias_k,
+          bias_v,
           add_zero_attn,
           dropout_p,
-          out_proj_weight, out_proj_bias,
+          out_proj_weight,
+          out_proj_bias,
           training: true,
           key_padding_mask: nil,
           need_weights: true,
           attn_mask: nil,
           use_separate_proj_weight: false,
-          q_proj_weight: nil, k_proj_weight: nil, v_proj_weight: nil,
-          static_k: nil, static_v: nil
+          q_proj_weight: nil,
+          k_proj_weight: nil,
+          v_proj_weight: nil,
+          static_k: nil,
+          static_v: nil
         )
-
           tgt_len, bsz, embed_dim = query.shape
           src_len = key.shape.first
 

data/lib/torch/nn/multihead_attention.rb

@@ -2,11 +2,18 @@ module Torch
   module NN
     class MultiheadAttention < Module
       def initialize(
-        embed_dim, num_heads,
-        dropout: 0.0, bias: true, add_bias_kv: false, add_zero_attn: false,
-        kdim: nil, vdim: nil, batch_first: false, device: nil, dtype: nil
+        embed_dim,
+        num_heads,
+        dropout: 0.0,
+        bias: true,
+        add_bias_kv: false,
+        add_zero_attn: false,
+        kdim: nil,
+        vdim: nil,
+        batch_first: false,
+        device: nil,
+        dtype: nil
       )
-
         super()
 
         @embed_dim = embed_dim
@@ -77,10 +84,13 @@ module Torch
       end
 
       def forward(
-        query, key, value,
-        key_padding_mask: nil, need_weights: true, attn_mask: nil
+        query,
+        key,
+        value,
+        key_padding_mask: nil,
+        need_weights: true,
+        attn_mask: nil
       )
-
         if batch_first?
           query, key, value = [query, key, value].map { |t| t.transpose(1, 0) }
         end

data/lib/torch/nn/rnn_base.rb

@@ -1,9 +1,16 @@
 module Torch
   module NN
     class RNNBase < Module
-      def initialize(mode, input_size, hidden_size, num_layers: 1, bias: true,
-        batch_first: false, dropout: 0.0, bidirectional: false)
-
+      def initialize(
+        mode,
+        input_size,
+        hidden_size,
+        num_layers: 1,
+        bias: true,
+        batch_first: false,
+        dropout: 0.0,
+        bidirectional: false
+      )
         super()
         @mode = mode
         @input_size = input_size

data/lib/torch/nn/transformer.rb

@@ -7,13 +7,18 @@ module Torch
   module NN
     class Transformer < Module
       def initialize(
-        d_model: 512, nhead: 8,
-        num_encoder_layers: 6, num_decoder_layers: 6,
-        dim_feedforward: 2048, dropout: 0.1, activation: :relu,
-        custom_encoder: nil, custom_decoder: nil,
-        layer_norm_eps: 1e-5, batch_first: false
+        d_model: 512,
+        nhead: 8,
+        num_encoder_layers: 6,
+        num_decoder_layers: 6,
+        dim_feedforward: 2048,
+        dropout: 0.1,
+        activation: :relu,
+        custom_encoder: nil,
+        custom_decoder: nil,
+        layer_norm_eps: 1e-5,
+        batch_first: false
       )
-
         super()
 
         @encoder =
@@ -60,11 +65,15 @@ module Torch
       end
 
       def forward(
-        src, tgt,
-        src_mask: nil, tgt_mask: nil, memory_mask: nil,
-        src_key_padding_mask: nil, tgt_key_padding_mask: nil, memory_key_padding_mask: nil
+        src,
+        tgt,
+        src_mask: nil,
+        tgt_mask: nil,
+        memory_mask: nil,
+        src_key_padding_mask: nil,
+        tgt_key_padding_mask: nil,
+        memory_key_padding_mask: nil
       )
-
         if (!batch_first? && src.size(1) != tgt.size(1)) ||
           (batch_first? && src.size(0) != tgt.size(0))