torch-rb 0.3.1 → 0.3.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 06e94b492acbbdb71f9e6a11081fb043a03ae0d5c704cc79faa31dd96bde70ef
-  data.tar.gz: 4f38fa52d30ef9bf121204423b4d675f21dbef806b6f137152f2cf9399ddf4bb
+  metadata.gz: f7b85027dfbb5a3d8de3741d00f4256fd13a6b5496a123564472a48d8a084c1b
+  data.tar.gz: 5c684e45ec115ce3b9cc5a3e223ee73cac3dce5fbacae1bd4d4faa7cf49adc5f
 SHA512:
-  metadata.gz: 2fb2613ca629a70f55009b697b15830d59c0d8fc06c1c5102917b4870cb783427fb56ecc08889c09e15c342381385f258b2a33102dc5adddf2d463d41674994d
-  data.tar.gz: f26a6ba91caa57a92b8b047217a35c39d1e9c4c361df77e2182053b4ab490f20792fc88dba169dae87d4a3d4ee4d69e2c779efb1fa6150b4d3f0d93e3762aec9
+  metadata.gz: acb727d9836709e5db4df21aeb6eec401e10f3c1910f95877493a9b1920cef4a0bd4914b906dab9b2ec18071fe95bf50de91a8a00a0914f3876ecb851e1c19c7
+  data.tar.gz: 7e69bde091825d7dcda81cfcfebd220c8072442322071a73eafd2849d9a899229bbcc2ce2b80f78e4b44e468e7a263ec83fd9a3fb3c1bf3073573596f40ec143

data/CHANGELOG.md

@@ -1,3 +1,31 @@
+## 0.3.6 (2020-09-17)
+
+- Added `inplace` option for leaky ReLU
+- Fixed error with methods that return a tensor list (`chunk`, `split`, and `unbind`)
+- Fixed error with buffers on GPU
+
+## 0.3.5 (2020-09-04)
+
+- Fixed error with data loader (due to `dtype` of `randperm`)
+
+## 0.3.4 (2020-08-26)
+
+- Added `Torch.clamp` method
+
+## 0.3.3 (2020-08-25)
+
+- Added spectral ops
+- Fixed tensor indexing
+
+## 0.3.2 (2020-08-24)
+
+- Added `enable_grad` method
+- Added `random_split` method
+- Added `collate_fn` option to `DataLoader`
+- Added `grad=` method to `Tensor`
+- Fixed error with `grad` method when empty
+- Fixed `EmbeddingBag`
+
 ## 0.3.1 (2020-08-17)
 
 - Added `create_graph` and `retain_graph` options to `backward` method

data/README.md

@@ -2,7 +2,11 @@
 
 :fire: Deep learning for Ruby, powered by [LibTorch](https://pytorch.org)
 
-For computer vision tasks, also check out [TorchVision](https://github.com/ankane/torchvision)
+Check out:
+
+- [TorchVision](https://github.com/ankane/torchvision) for computer vision tasks
+- [TorchText](https://github.com/ankane/torchtext) for text and NLP tasks
+- [TorchAudio](https://github.com/ankane/torchaudio) for audio tasks
 
 [![Build Status](https://travis-ci.org/ankane/torch.rb.svg?branch=master)](https://travis-ci.org/ankane/torch.rb)
 
@@ -411,7 +415,7 @@ Here’s the list of compatible versions.
 
 Torch.rb | LibTorch
 --- | ---
-0.3.0-0.3.1 | 1.6.0
+0.3.0-0.3.4 | 1.6.0
 0.2.0-0.2.7 | 1.5.0-1.5.1
 0.1.8 | 1.4.0
 0.1.0-0.1.7 | 1.3.1

data/ext/torch/ext.cpp

@@ -16,6 +16,7 @@
 #include "nn_functions.hpp"
 
 using namespace Rice;
+using torch::indexing::TensorIndex;
 
 // need to make a distinction between parameters and tensors
 class Parameter: public torch::autograd::Variable {
@@ -28,6 +29,15 @@ void handle_error(torch::Error const & ex)
   throw Exception(rb_eRuntimeError, ex.what_without_backtrace());
 }
 
+std::vector<TensorIndex> index_vector(Array a) {
+  auto indices = std::vector<TensorIndex>();
+  indices.reserve(a.size());
+  for (size_t i = 0; i < a.size(); i++) {
+    indices.push_back(from_ruby<TensorIndex>(a[i]));
+  }
+  return indices;
+}
+
 extern "C"
 void Init_ext()
 {
@@ -58,6 +68,13 @@ void Init_ext()
         return generator.seed();
       });
 
+  Class rb_cTensorIndex = define_class_under<TensorIndex>(rb_mTorch, "TensorIndex")
+    .define_singleton_method("boolean", *[](bool value) { return TensorIndex(value); })
+    .define_singleton_method("integer", *[](int64_t value) { return TensorIndex(value); })
+    .define_singleton_method("tensor", *[](torch::Tensor& value) { return TensorIndex(value); })
+    .define_singleton_method("slice", *[](torch::optional<int64_t> start_index, torch::optional<int64_t> stop_index) { return TensorIndex(torch::indexing::Slice(start_index, stop_index)); })
+    .define_singleton_method("none", *[]() { return TensorIndex(torch::indexing::None); });
+
   // https://pytorch.org/cppdocs/api/structc10_1_1_i_value.html
   Class rb_cIValue = define_class_under<torch::IValue>(rb_mTorch, "IValue")
     .add_handler<torch::Error>(handle_error)
@@ -284,11 +301,6 @@ void Init_ext()
         // https://github.com/pytorch/pytorch/issues/20356#issuecomment-567663701
         return torch::pickle_load(v);
       })
-    .define_singleton_method(
-      "_binary_cross_entropy_with_logits",
-      *[](const Tensor &input, const Tensor &target, OptionalTensor weight, OptionalTensor pos_weight, MyReduction reduction) {
-        return torch::binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction);
-      })
     .define_singleton_method(
       "_from_blob",
       *[](String s, IntArrayRef size, const torch::TensorOptions &options) {
@@ -330,6 +342,18 @@ void Init_ext()
     .define_method("numel", &torch::Tensor::numel)
     .define_method("element_size", &torch::Tensor::element_size)
     .define_method("requires_grad", &torch::Tensor::requires_grad)
+    .define_method(
+      "_index",
+      *[](Tensor& self, Array indices) {
+        auto vec = index_vector(indices);
+        return self.index(vec);
+      })
+    .define_method(
+      "_index_put_custom",
+      *[](Tensor& self, Array indices, torch::Tensor& value) {
+        auto vec = index_vector(indices);
+        return self.index_put_(vec, value);
+      })
     .define_method(
       "contiguous?",
       *[](Tensor& self) {
@@ -350,15 +374,16 @@ void Init_ext()
       *[](Tensor& self, bool requires_grad) {
         return self.set_requires_grad(requires_grad);
       })
-    .define_method(
-      "_backward",
-      *[](Tensor& self, OptionalTensor gradient, bool create_graph, bool retain_graph) {
-        return self.backward(gradient, create_graph, retain_graph);
-      })
     .define_method(
       "grad",
       *[](Tensor& self) {
-        return self.grad();
+        auto grad = self.grad();
+        return grad.defined() ? to_ruby<torch::Tensor>(grad) : Nil;
+      })
+    .define_method(
+      "grad=",
+      *[](Tensor& self, torch::Tensor& grad) {
+        self.grad() = grad;
       })
     .define_method(
       "_dtype",
@@ -502,6 +527,7 @@ void Init_ext()
       });
 
   Module rb_mInit = define_module_under(rb_mNN, "Init")
+    .add_handler<torch::Error>(handle_error)
     .define_singleton_method(
       "_calculate_gain",
       *[](NonlinearityType nonlinearity, double param) {
@@ -580,11 +606,16 @@ void Init_ext()
       *[](Parameter& self) {
         auto grad = self.grad();
         return grad.defined() ? to_ruby<torch::Tensor>(grad) : Nil;
+      })
+    .define_method(
+      "grad=",
+      *[](Parameter& self, torch::Tensor& grad) {
+        self.grad() = grad;
       });
 
   Class rb_cDevice = define_class_under<torch::Device>(rb_mTorch, "Device")
-    .define_constructor(Constructor<torch::Device, std::string>())
     .add_handler<torch::Error>(handle_error)
+    .define_constructor(Constructor<torch::Device, std::string>())
     .define_method("index", &torch::Device::index)
     .define_method("index?", &torch::Device::has_index)
     .define_method(

data/ext/torch/templates.cpp

@@ -53,3 +53,11 @@ Object wrap(std::tuple<torch::Tensor, torch::Tensor, double, int64_t> x) {
   a.push(to_ruby<int64_t>(std::get<3>(x)));
   return Object(a);
 }
+
+Object wrap(std::vector<torch::Tensor> x) {
+  Array a;
+  for (auto& t : x) {
+    a.push(to_ruby<torch::Tensor>(t));
+  }
+  return Object(a);
+}

data/ext/torch/templates.hpp

@@ -9,6 +9,11 @@
 
 using namespace Rice;
 
+using torch::Device;
+using torch::Scalar;
+using torch::ScalarType;
+using torch::Tensor;
+
 // need to wrap torch::IntArrayRef() since
 // it doesn't own underlying data
 class IntArrayRef {
@@ -32,30 +37,6 @@ IntArrayRef from_ruby<IntArrayRef>(Object x)
   return IntArrayRef(x);
 }
 
-// for now
-class Scalar {
-  torch::Scalar value;
-  public:
-    Scalar(Object o) {
-      // TODO cast based on Ruby type
-      if (o.rb_type() == T_FIXNUM) {
-        value = torch::Scalar(from_ruby<int64_t>(o));
-      } else {
-        value = torch::Scalar(from_ruby<float>(o));
-      }
-    }
-    operator torch::Scalar() {
-      return value;
-    }
-};
-
-template<>
-inline
-Scalar from_ruby<Scalar>(Object x)
-{
-  return Scalar(x);
-}
-
 class TensorList {
   std::vector<torch::Tensor> vec;
   public:
@@ -174,8 +155,6 @@ MyReduction from_ruby<MyReduction>(Object x)
   return MyReduction(x);
 }
 
-typedef torch::Tensor Tensor;
-
 class OptionalTensor {
   Object value;
   public:
@@ -190,6 +169,17 @@ class OptionalTensor {
     }
 };
 
+template<>
+inline
+Scalar from_ruby<Scalar>(Object x)
+{
+  if (x.rb_type() == T_FIXNUM) {
+    return torch::Scalar(from_ruby<int64_t>(x));
+  } else {
+    return torch::Scalar(from_ruby<double>(x));
+  }
+}
+
 template<>
 inline
 OptionalTensor from_ruby<OptionalTensor>(Object x)
@@ -197,46 +187,60 @@ OptionalTensor from_ruby<OptionalTensor>(Object x)
   return OptionalTensor(x);
 }
 
-class ScalarType {
-  Object value;
-  public:
-    ScalarType(Object o) {
-      value = o;
-    }
-    operator at::ScalarType() {
-      throw std::runtime_error("ScalarType arguments not implemented yet");
-    }
-};
+template<>
+inline
+torch::optional<torch::ScalarType> from_ruby<torch::optional<torch::ScalarType>>(Object x)
+{
+  if (x.is_nil()) {
+    return torch::nullopt;
+  } else {
+    return torch::optional<torch::ScalarType>{from_ruby<torch::ScalarType>(x)};
+  }
+}
 
 template<>
 inline
-ScalarType from_ruby<ScalarType>(Object x)
+torch::optional<int64_t> from_ruby<torch::optional<int64_t>>(Object x)
 {
-  return ScalarType(x);
+  if (x.is_nil()) {
+    return torch::nullopt;
+  } else {
+    return torch::optional<int64_t>{from_ruby<int64_t>(x)};
+  }
 }
 
-class OptionalScalarType {
-  Object value;
-  public:
-    OptionalScalarType(Object o) {
-      value = o;
-    }
-    operator c10::optional<at::ScalarType>() {
-      if (value.is_nil()) {
-        return c10::nullopt;
-      }
-      return ScalarType(value);
-    }
-};
+template<>
+inline
+torch::optional<double> from_ruby<torch::optional<double>>(Object x)
+{
+  if (x.is_nil()) {
+    return torch::nullopt;
+  } else {
+    return torch::optional<double>{from_ruby<double>(x)};
+  }
+}
 
 template<>
 inline
-OptionalScalarType from_ruby<OptionalScalarType>(Object x)
+torch::optional<bool> from_ruby<torch::optional<bool>>(Object x)
 {
-  return OptionalScalarType(x);
+  if (x.is_nil()) {
+    return torch::nullopt;
+  } else {
+    return torch::optional<bool>{from_ruby<bool>(x)};
+  }
 }
 
-typedef torch::Device Device;
+template<>
+inline
+torch::optional<Scalar> from_ruby<torch::optional<Scalar>>(Object x)
+{
+  if (x.is_nil()) {
+    return torch::nullopt;
+  } else {
+    return torch::optional<Scalar>{from_ruby<Scalar>(x)};
+  }
+}
 
 Object wrap(std::tuple<torch::Tensor, torch::Tensor> x);
 Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> x);
@@ -244,3 +248,4 @@ Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tenso
 Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> x);
 Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, int64_t> x);
 Object wrap(std::tuple<torch::Tensor, torch::Tensor, double, int64_t> x);
+Object wrap(std::vector<torch::Tensor> x);

data/lib/torch.rb

@@ -179,8 +179,10 @@ require "torch/nn/functional"
 require "torch/nn/init"
 
 # utils
+require "torch/utils/data"
 require "torch/utils/data/data_loader"
 require "torch/utils/data/dataset"
+require "torch/utils/data/subset"
 require "torch/utils/data/tensor_dataset"
 
 # hub
@@ -237,25 +239,22 @@ module Torch
     cls
   end
 
-  FloatTensor = _make_tensor_class(:float32)
-  DoubleTensor = _make_tensor_class(:float64)
-  HalfTensor = _make_tensor_class(:float16)
-  ByteTensor = _make_tensor_class(:uint8)
-  CharTensor = _make_tensor_class(:int8)
-  ShortTensor = _make_tensor_class(:int16)
-  IntTensor = _make_tensor_class(:int32)
-  LongTensor = _make_tensor_class(:int64)
-  BoolTensor = _make_tensor_class(:bool)
-
-  CUDA::FloatTensor = _make_tensor_class(:float32, true)
-  CUDA::DoubleTensor = _make_tensor_class(:float64, true)
-  CUDA::HalfTensor = _make_tensor_class(:float16, true)
-  CUDA::ByteTensor = _make_tensor_class(:uint8, true)
-  CUDA::CharTensor = _make_tensor_class(:int8, true)
-  CUDA::ShortTensor = _make_tensor_class(:int16, true)
-  CUDA::IntTensor = _make_tensor_class(:int32, true)
-  CUDA::LongTensor = _make_tensor_class(:int64, true)
-  CUDA::BoolTensor = _make_tensor_class(:bool, true)
+  DTYPE_TO_CLASS = {
+    float32: "FloatTensor",
+    float64: "DoubleTensor",
+    float16: "HalfTensor",
+    uint8: "ByteTensor",
+    int8: "CharTensor",
+    int16: "ShortTensor",
+    int32: "IntTensor",
+    int64: "LongTensor",
+    bool: "BoolTensor"
+  }
+
+  DTYPE_TO_CLASS.each do |dtype, class_name|
+    const_set(class_name, _make_tensor_class(dtype))
+    CUDA.const_set(class_name, _make_tensor_class(dtype, true))
+  end
 
   class << self
     # Torch.float, Torch.long, etc
@@ -316,6 +315,16 @@ module Torch
       end
     end
 
+    def enable_grad
+      previous_value = grad_enabled?
+      begin
+        _set_grad_enabled(true)
+        yield
+      ensure
+        _set_grad_enabled(previous_value)
+      end
+    end
+
     def device(str)
       Device.new(str)
     end
@@ -376,6 +385,10 @@ module Torch
     end
 
     def randperm(n, **options)
+      # dtype hack in Python
+      # https://github.com/pytorch/pytorch/blob/v1.6.0/tools/autograd/gen_python_functions.py#L1307-L1311
+      options[:dtype] ||= :int64
+
       _randperm(n, tensor_options(**options))
     end
 
@@ -448,6 +461,22 @@ module Torch
       zeros(input.size, **like_options(input, options))
     end
 
+    def stft(input, n_fft, hop_length: nil, win_length: nil, window: nil, center: true, pad_mode: "reflect", normalized: false, onesided: true)
+      if center
+        signal_dim = input.dim
+        extended_shape = [1] * (3 - signal_dim) + input.size
+        pad = n_fft.div(2).to_i
+        input = NN::F.pad(input.view(extended_shape), [pad, pad], mode: pad_mode)
+        input = input.view(input.shape[-signal_dim..-1])
+      end
+      _stft(input, n_fft, hop_length, win_length, window, normalized, onesided)
+    end
+
+    def clamp(tensor, min, max)
+      tensor = _clamp_min(tensor, min)
+      _clamp_max(tensor, max)
+    end
+
     private
 
     def to_ivalue(obj)

data/lib/torch/hub.rb

@@ -7,25 +7,26 @@ module Torch
 
       def download_url_to_file(url, dst)
         uri = URI(url)
-        tmp = "#{Dir.tmpdir}/#{Time.now.to_f}" # TODO better name
+        tmp = nil
         location = nil
 
+        puts "Downloading #{url}..."
         Net::HTTP.start(uri.host, uri.port, use_ssl: uri.scheme == "https") do |http|
           request = Net::HTTP::Get.new(uri)
 
-          puts "Downloading #{url}..."
-          File.open(tmp, "wb") do |f|
-            http.request(request) do |response|
-              case response
-              when Net::HTTPRedirection
-                location = response["location"]
-              when Net::HTTPSuccess
+          http.request(request) do |response|
+            case response
+            when Net::HTTPRedirection
+              location = response["location"]
+            when Net::HTTPSuccess
+              tmp = "#{Dir.tmpdir}/#{Time.now.to_f}" # TODO better name
+              File.open(tmp, "wb") do |f|
                 response.read_body do |chunk|
                   f.write(chunk)
                 end
-              else
-                raise Error, "Bad response"
               end
+            else
+              raise Error, "Bad response"
             end
           end
         end

data/lib/torch/native/function.rb

@@ -1,10 +1,14 @@
 module Torch
   module Native
     class Function
-      attr_reader :function
+      attr_reader :function, :tensor_options
 
       def initialize(function)
         @function = function
+
+        tensor_options_str = ", *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None)"
+        @tensor_options = @function["func"].include?(tensor_options_str)
+        @function["func"].sub!(tensor_options_str, ")")
       end
 
       def func
@@ -82,6 +86,10 @@ module Torch
         @ret_size ||= func.split("->").last.split(", ").size
       end
 
+      def ret_array?
+        @ret_array ||= func.split("->").last.include?('[]')
+      end
+
       def out?
         out_size > 0 && base_name[-1] != "_"
       end

data/lib/torch/native/generator.rb

@@ -18,12 +18,12 @@ module Torch
           functions = functions()
 
           # skip functions
-          skip_args = ["bool[3]", "Dimname", "Layout", "Storage", "ConstQuantizerPtr"]
+          skip_args = ["Layout", "Storage", "ConstQuantizerPtr"]
 
           # remove functions
           functions.reject! do |f|
             f.ruby_name.start_with?("_") ||
-            f.ruby_name.end_with?("_backward") ||
+            f.ruby_name.include?("_backward") ||
             f.args.any? { |a| a[:type].include?("Dimname") }
           end
 
@@ -31,32 +31,15 @@ module Torch
           todo_functions, functions =
             functions.partition do |f|
               f.args.any? do |a|
-                a[:type].include?("?") && !["Tensor?", "Generator?", "int?", "ScalarType?", "Tensor?[]"].include?(a[:type]) ||
                 skip_args.any? { |sa| a[:type].include?(sa) } ||
+                # call to 'range' is ambiguous
+                f.cpp_name == "_range" ||
                 # native_functions.yaml is missing size argument for normal
                 # https://pytorch.org/cppdocs/api/function_namespacetorch_1a80253fe5a3ded4716ec929a348adb4b9.html
                 (f.base_name == "normal" && !f.out?)
               end
             end
 
-          # generate additional functions for optional arguments
-          # there may be a better way to do this
-          optional_functions, functions = functions.partition { |f| f.args.any? { |a| a[:type] == "int?" } }
-          optional_functions.each do |f|
-            next if f.ruby_name == "cross"
-            next if f.ruby_name.start_with?("avg_pool") && f.out?
-
-            opt_args = f.args.select { |a| a[:type] == "int?" }
-            if opt_args.size == 1
-              sep = f.name.include?(".") ? "_" : "."
-              f1 = Function.new(f.function.merge("func" => f.func.sub("(", "#{sep}#{opt_args.first[:name]}(").gsub("int?", "int")))
-              # TODO only remove some arguments
-              f2 = Function.new(f.function.merge("func" => f.func.sub(/, int\?.+\) ->/, ") ->")))
-              functions << f1
-              functions << f2
-            end
-          end
-
           # todo_functions.each do |f|
           #   puts f.func
           #   puts
@@ -97,7 +80,8 @@ void add_%{type}_functions(Module m) {
 
           cpp_defs = []
           functions.sort_by(&:cpp_name).each do |func|
-            fargs = func.args #.select { |a| a[:type] != "Generator?" }
+            fargs = func.args.dup #.select { |a| a[:type] != "Generator?" }
+            fargs << {name: "options", type: "TensorOptions"} if func.tensor_options
 
             cpp_args = []
             fargs.each do |a|
@@ -109,7 +93,7 @@ void add_%{type}_functions(Module m) {
                   # TODO better signature
                   "OptionalTensor"
                 when "ScalarType?"
-                  "OptionalScalarType"
+                  "torch::optional<ScalarType>"
                 when "Tensor[]"
                   "TensorList"
                 when "Tensor?[]"
@@ -117,6 +101,14 @@ void add_%{type}_functions(Module m) {
                   "TensorList"
                 when "int"
                   "int64_t"
+                when "int?"
+                  "torch::optional<int64_t>"
+                when "float?"
+                  "torch::optional<double>"
+                when "bool?"
+                  "torch::optional<bool>"
+                when "Scalar?"
+                  "torch::optional<torch::Scalar>"
                 when "float"
                   "double"
                 when /\Aint\[/
@@ -125,6 +117,8 @@ void add_%{type}_functions(Module m) {
                   "Tensor &"
                 when "str"
                   "std::string"
+                when "TensorOptions"
+                  "const torch::TensorOptions &"
                 else
                   a[:type]
                 end
@@ -141,8 +135,8 @@ void add_%{type}_functions(Module m) {
             prefix = def_method == :define_method ? "self." : "torch::"
 
             body = "#{prefix}#{dispatch}(#{args.join(", ")})"
-            # TODO check type as well
-            if func.ret_size > 1
+
+            if func.ret_size > 1 || func.ret_array?
               body = "wrap(#{body})"
             end
 

data/lib/torch/native/parser.rb

@@ -83,6 +83,12 @@ module Torch
                 else
                   v.is_a?(Integer)
                 end
+              when "int?"
+                v.is_a?(Integer) || v.nil?
+              when "float?"
+                v.is_a?(Numeric) || v.nil?
+              when "bool?"
+                v == true || v == false || v.nil?
               when "float"
                 v.is_a?(Numeric)
               when /int\[.*\]/
@@ -95,6 +101,10 @@ module Torch
                 v.is_a?(Array) && v.all? { |v2| v2.is_a?(Integer) }
               when "Scalar"
                 v.is_a?(Numeric)
+              when "Scalar?"
+                v.is_a?(Numeric) || v.nil?
+              when "ScalarType"
+                false # not supported yet
               when "ScalarType?"
                 v.nil?
               when "bool"
@@ -126,9 +136,11 @@ module Torch
         end
 
         func = candidates.first
+        args = func.args.map { |a| final_values[a[:name]] }
+        args << TensorOptions.new.dtype(6) if func.tensor_options
         {
           name: func.cpp_name,
-          args: func.args.map { |a| final_values[a[:name]] }
+          args: args
         }
       end
     end

data/lib/torch/nn/functional.rb

@@ -178,8 +178,12 @@ module Torch
           Torch.hardshrink(input, lambd)
         end
 
-        def leaky_relu(input, negative_slope = 0.01)
-          NN.leaky_relu(input, negative_slope)
+        def leaky_relu(input, negative_slope = 0.01, inplace: false)
+          if inplace
+            NN.leaky_relu!(input, negative_slope)
+          else
+            NN.leaky_relu(input, negative_slope)
+          end
         end
 
         def log_sigmoid(input)
@@ -373,7 +377,8 @@ module Torch
             end
 
           # weight and input swapped
-          Torch.embedding_bag(weight, input, offsets, scale_grad_by_freq, mode_enum, sparse, per_sample_weights)
+          ret, _, _, _ = Torch.embedding_bag(weight, input, offsets, scale_grad_by_freq, mode_enum, sparse, per_sample_weights)
+          ret
         end
 
         # distance functions
@@ -426,6 +431,9 @@ module Torch
         end
 
         def mse_loss(input, target, reduction: "mean")
+          if target.size != input.size
+            warn "Using a target size (#{target.size}) that is different to the input size (#{input.size}). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size."
+          end
           NN.mse_loss(input, target, reduction)
         end
 

data/lib/torch/nn/leaky_relu.rb

@@ -1,14 +1,14 @@
 module Torch
   module NN
     class LeakyReLU < Module
-      def initialize(negative_slope: 1e-2) #, inplace: false)
+      def initialize(negative_slope: 1e-2, inplace: false)
         super()
         @negative_slope = negative_slope
-        # @inplace = inplace
+        @inplace = inplace
       end
 
       def forward(input)
-        F.leaky_relu(input, @negative_slope) #, inplace: @inplace)
+        F.leaky_relu(input, @negative_slope, inplace: @inplace)
       end
 
       def extra_inspect

data/lib/torch/nn/module.rb

@@ -55,7 +55,12 @@ module Torch
             end
           end
         end
-        # TODO apply to more objects
+
+        @buffers.each_key do |k|
+          buf = @buffers[k]
+          @buffers[k] = fn.call(buf) unless buf.nil?
+        end
+
         self
       end
 

data/lib/torch/tensor.rb

@@ -103,11 +103,6 @@ module Torch
       Torch.empty(0, dtype: dtype)
     end
 
-    def backward(gradient = nil, retain_graph: nil, create_graph: false)
-      retain_graph = create_graph if retain_graph.nil?
-      _backward(gradient, retain_graph, create_graph)
-    end
-
     # TODO read directly from memory
     def numo
       cls = Torch._dtype_to_numo[dtype]
@@ -188,49 +183,15 @@ module Torch
     # based on python_variable_indexing.cpp and
     # https://pytorch.org/cppdocs/notes/tensor_indexing.html
     def [](*indexes)
-      result = self
-      dim = 0
-      indexes.each do |index|
-        if index.is_a?(Numeric)
-          result = result._select_int(dim, index)
-        elsif index.is_a?(Range)
-          finish = index.end
-          finish += 1 unless index.exclude_end?
-          result = result._slice_tensor(dim, index.begin, finish, 1)
-          dim += 1
-        elsif index.is_a?(Tensor)
-          result = result.index([index])
-        elsif index.nil?
-          result = result.unsqueeze(dim)
-          dim += 1
-        elsif index == true
-          result = result.unsqueeze(dim)
-          # TODO handle false
-        else
-          raise Error, "Unsupported index type: #{index.class.name}"
-        end
-      end
-      result
+      _index(tensor_indexes(indexes))
     end
 
     # based on python_variable_indexing.cpp and
     # https://pytorch.org/cppdocs/notes/tensor_indexing.html
-    def []=(index, value)
+    def []=(*indexes, value)
       raise ArgumentError, "Tensor does not support deleting items" if value.nil?
-
       value = Torch.tensor(value, dtype: dtype) unless value.is_a?(Tensor)
-
-      if index.is_a?(Numeric)
-        index_put!([Torch.tensor(index)], value)
-      elsif index.is_a?(Range)
-        finish = index.end
-        finish += 1 unless index.exclude_end?
-        _slice_tensor(0, index.begin, finish, 1).copy!(value)
-      elsif index.is_a?(Tensor)
-        index_put!([index], value)
-      else
-        raise Error, "Unsupported index type: #{index.class.name}"
-      end
+      _index_put_custom(tensor_indexes(indexes), value)
     end
 
     # native functions that need manually defined
@@ -244,13 +205,13 @@ module Torch
       end
     end
 
-    # native functions overlap, so need to handle manually
+    # parser can't handle overlap, so need to handle manually
     def random!(*args)
       case args.size
       when 1
         _random__to(*args)
       when 2
-        _random__from_to(*args)
+        _random__from(*args)
       else
         _random_(*args)
       end
@@ -260,5 +221,32 @@ module Torch
       _clamp_min_(min)
       _clamp_max_(max)
     end
+
+    private
+
+    def tensor_indexes(indexes)
+      indexes.map do |index|
+        case index
+        when Integer
+          TensorIndex.integer(index)
+        when Range
+          finish = index.end || -1
+          if finish == -1 && !index.exclude_end?
+            finish = nil
+          else
+            finish += 1 unless index.exclude_end?
+          end
+          TensorIndex.slice(index.begin, finish)
+        when Tensor
+          TensorIndex.tensor(index)
+        when nil
+          TensorIndex.none
+        when true, false
+          TensorIndex.boolean(index)
+        else
+          raise Error, "Unsupported index type: #{index.class.name}"
+        end
+      end
+    end
   end
 end

data/lib/torch/utils/data.rb

@@ -0,0 +1,23 @@
+module Torch
+  module Utils
+    module Data
+      class << self
+        def random_split(dataset, lengths)
+          if lengths.sum != dataset.length
+            raise ArgumentError, "Sum of input lengths does not equal the length of the input dataset!"
+          end
+
+          indices = Torch.randperm(lengths.sum).to_a
+          _accumulate(lengths).zip(lengths).map { |offset, length| Subset.new(dataset, indices[(offset - length)...offset]) }
+        end
+
+        private
+
+        def _accumulate(iterable)
+          sum = 0
+          iterable.map { |x| sum += x }
+        end
+      end
+    end
+  end
+end

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

@@ -6,10 +6,22 @@ module Torch
 
         attr_reader :dataset
 
-        def initialize(dataset, batch_size: 1, shuffle: false)
+        def initialize(dataset, batch_size: 1, shuffle: false, collate_fn: nil)
           @dataset = dataset
           @batch_size = batch_size
           @shuffle = shuffle
+
+          @batch_sampler = nil
+
+          if collate_fn.nil?
+            if auto_collation?
+              collate_fn = method(:default_collate)
+            else
+              collate_fn = method(:default_convert)
+            end
+          end
+
+          @collate_fn = collate_fn
         end
 
         def each
@@ -25,18 +37,20 @@ module Torch
             end
 
           indexes.each_slice(@batch_size) do |idx|
-            batch = idx.map { |i| @dataset[i] }
-            yield collate(batch)
+            # TODO improve performance
+            yield @collate_fn.call(idx.map { |i| @dataset[i] })
           end
         end
 
         def size
           (@dataset.size / @batch_size.to_f).ceil
         end
+        alias_method :length, :size
+        alias_method :count, :size
 
         private
 
-        def collate(batch)
+        def default_convert(batch)
           elem = batch[0]
           case elem
           when Tensor
@@ -44,11 +58,15 @@ module Torch
           when Integer
             Torch.tensor(batch)
           when Array
-            batch.transpose.map { |v| collate(v) }
+            batch.transpose.map { |v| default_convert(v) }
           else
-            raise NotImpelmentYet
+            raise NotImplementedYet
           end
         end
+
+        def auto_collation?
+          !@batch_sampler.nil?
+        end
       end
     end
   end

data/lib/torch/utils/data/subset.rb

@@ -0,0 +1,25 @@
+module Torch
+  module Utils
+    module Data
+      class Subset < Dataset
+        def initialize(dataset, indices)
+          @dataset = dataset
+          @indices = indices
+        end
+
+        def [](idx)
+          @dataset[@indices[idx]]
+        end
+
+        def length
+          @indices.length
+        end
+        alias_method :size, :length
+
+        def to_a
+          @indices.map { |i| @dataset[i] }
+        end
+      end
+    end
+  end
+end

data/lib/torch/utils/data/tensor_dataset.rb

@@ -16,6 +16,8 @@ module Torch
         def size
           @tensors[0].size(0)
         end
+        alias_method :length, :size
+        alias_method :count, :size
       end
     end
   end

data/lib/torch/version.rb

@@ -1,3 +1,3 @@
 module Torch
-  VERSION = "0.3.1"
+  VERSION = "0.3.6"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: torch-rb
 version: !ruby/object:Gem::Version
-  version: 0.3.1
+  version: 0.3.6
 platform: ruby
 authors:
 - Andrew Kane
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2020-08-17 00:00:00.000000000 Z
+date: 2020-09-18 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rice
@@ -259,8 +259,10 @@ files:
 - lib/torch/optim/rprop.rb
 - lib/torch/optim/sgd.rb
 - lib/torch/tensor.rb
+- lib/torch/utils/data.rb
 - lib/torch/utils/data/data_loader.rb
 - lib/torch/utils/data/dataset.rb
+- lib/torch/utils/data/subset.rb
 - lib/torch/utils/data/tensor_dataset.rb
 - lib/torch/version.rb
 homepage: https://github.com/ankane/torch.rb