torch-rb 0.3.2 → 0.3.7

data/lib/torch/native/parser.rb

@@ -6,14 +6,24 @@ module Torch
         @name = @functions.first.ruby_name
         @min_args = @functions.map { |f| f.args.count { |a| a[:pos] && !a[:has_default] } }.min
         @max_args = @functions.map { |f| f.args.count { |a| a[:pos] } }.max
+        @int_array_first = @functions.all? { |c| c.args.first && c.args.first[:type] == "int[]" }
       end
 
+      # TODO improve performance
+      # possibly move to C++ (see python_arg_parser.cpp)
       def parse(args, options)
         candidates = @functions.dup
 
-        # remove nil
-        while args.any? && args.last.nil?
-          args.pop
+        # TODO check candidates individually to see if they match
+        if @int_array_first
+          int_args = []
+          while args.first.is_a?(Integer)
+            int_args << args.shift
+          end
+          if int_args.any?
+            raise ArgumentError, "argument '#{candidates.first.args.first[:name]}' must be array of ints, but found element of type #{args.first.class.name} at pos #{int_args.size + 1}" if args.any?
+            args.unshift(int_args)
+          end
         end
 
         # TODO account for args passed as options here
@@ -25,89 +35,60 @@ module Torch
 
         candidates.reject! { |f| args.size > f.args.size }
 
-        # exclude functions missing required options
-        candidates.reject! do |func|
-          # TODO make more generic
-          func.out? && !options[:out]
-        end
-
         # handle out with multiple
         # there should only be one match, so safe to modify all
-        out_func = candidates.find { |f| f.out? }
-        if out_func && out_func.out_size > 1 && options[:out]
-          out_args = out_func.args.last(2).map { |a| a[:name] }
-          out_args.zip(options.delete(:out)).each do |k, v|
-            options[k.to_sym] = v
-          end
-          candidates = [out_func]
-        end
-
-        # exclude functions where options don't match
-        options.each do |k, v|
-          candidates.select! do |func|
-            func.args.any? { |a| a[:name] == k.to_s }
+        if options[:out]
+          if (out_func = candidates.find { |f| f.out? }) && out_func.out_size > 1
+            out_args = out_func.args.last(2).map { |a| a[:name] }
+            out_args.zip(options.delete(:out)).each do |k, v|
+              options[k] = v
+            end
+            candidates = [out_func]
           end
-          # TODO show all bad keywords at once like Ruby?
-          return {error: "unknown keyword: #{k}"} if candidates.empty?
+        else
+          # exclude functions missing required options
+          candidates.reject!(&:out?)
         end
 
-        final_values = {}
+        final_values = nil
 
         # check args
-        candidates.select! do |func|
+        while (func = candidates.shift)
           good = true
 
-          values = args.zip(func.args).map { |a, fa| [fa[:name], a] }.to_h
-          values.merge!(options.map { |k, v| [k.to_s, v] }.to_h)
-          func.args.each do |fa|
-            values[fa[:name]] = fa[:default] if values[fa[:name]].nil?
+          # set values
+          # TODO use array instead of hash?
+          values = {}
+          args.each_with_index do |a, i|
+            values[func.arg_names[i]] = a
+          end
+          options.each do |k, v|
+            values[k] = v
+          end
+          func.arg_defaults.each do |k, v|
+            values[k] = v unless values.key?(k)
+          end
+          func.int_array_lengths.each do |k, len|
+            values[k] = [values[k]] * len if values[k].is_a?(Integer)
           end
 
-          arg_types = func.args.map { |a| [a[:name], a[:type]] }.to_h
+          arg_checkers = func.arg_checkers
 
           values.each_key do |k|
-            v = values[k]
-            t = arg_types[k].split("(").first
-
-            good =
-              case t
-              when "Tensor"
-                v.is_a?(Tensor)
-              when "Tensor?"
-                v.nil? || v.is_a?(Tensor)
-              when "Tensor[]", "Tensor?[]"
-                v.is_a?(Array) && v.all? { |v2| v2.is_a?(Tensor) }
-              when "int"
-                if k == "reduction"
-                  v.is_a?(String)
-                else
-                  v.is_a?(Integer)
-                end
-              when "float"
-                v.is_a?(Numeric)
-              when /int\[.*\]/
-                if v.is_a?(Integer)
-                  size = t[4..-2]
-                  raise Error, "Unknown size: #{size}. Please report a bug with #{@name}." unless size =~ /\A\d+\z/
-                  v = [v] * size.to_i
-                  values[k] = v
-                end
-                v.is_a?(Array) && v.all? { |v2| v2.is_a?(Integer) }
-              when "Scalar"
-                v.is_a?(Numeric)
-              when "ScalarType?"
-                v.nil?
-              when "bool"
-                v == true || v == false
-              when "str"
-                v.is_a?(String)
-              else
-                raise Error, "Unknown argument type: #{arg_types[k]}. Please report a bug with #{@name}."
+            unless arg_checkers.key?(k)
+              good = false
+              if candidates.empty?
+                # TODO show all bad keywords at once like Ruby?
+                return {error: "unknown keyword: #{k}"}
               end
+              break
+            end
 
-            if !good
-              if candidates.size == 1
-                k = "input" if k == "self"
+            unless arg_checkers[k].call(values[k])
+              good = false
+              if candidates.empty?
+                t = func.arg_types[k]
+                k = :input if k == :self
                 return {error: "#{@name}(): argument '#{k}' must be #{t}"}
               end
               break
@@ -116,19 +97,19 @@ module Torch
 
           if good
             final_values = values
+            break
           end
-
-          good
         end
 
-        if candidates.size != 1
+        unless final_values
           raise Error, "This should never happen. Please report a bug with #{@name}."
         end
 
-        func = candidates.first
+        args = func.arg_names.map { |k| final_values[k] }
+        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)
@@ -465,6 +469,77 @@ module Torch
           Torch.triplet_margin_loss(anchor, positive, negative, margin, p, eps, swap, reduction)
         end
 
+        # vision
+
+        def interpolate(input, size: nil, scale_factor: nil, mode: "nearest", align_corners: nil, recompute_scale_factor: nil)
+          if ["nearest", "area"].include?(mode)
+            unless align_corners.nil?
+              raise ArgumentError, "align_corners option can only be set with the interpolating modes: linear | bilinear | bicubic | trilinear"
+            end
+          else
+            if align_corners.nil?
+              align_corners = false
+            end
+          end
+
+          scale_factor_len = input.dim - 2
+          scale_factor_list = [nil] * scale_factor_len
+          # default value of recompute_scale_factor is False
+          if !scale_factor.nil? && (recompute_scale_factor == false || recompute_scale_factor.nil?)
+            if scale_factor.is_a?(Array)
+              _scale_factor_repeated = scale_factor
+            else
+              _scale_factor_repeated = [scale_factor] * scale_factor_len
+            end
+            scale_factor_list = _scale_factor_repeated
+          end
+
+          # Give this variable a short name because it has to be repeated multiple times below.
+          sfl = scale_factor_list
+
+          closed_over_args = [input, size, scale_factor, recompute_scale_factor]
+          output_size = _interp_output_size(closed_over_args)
+          if input.dim == 3 && mode == "nearest"
+            NN.upsample_nearest1d(input, output_size, sfl[0])
+          elsif input.dim == 4 && mode == "nearest"
+            NN.upsample_nearest2d(input, output_size, sfl[0], sfl[1])
+          elsif input.dim == 5 && mode == "nearest"
+            NN.upsample_nearest3d(input, output_size, sfl[0], sfl[1], sfl[2])
+          elsif input.dim == 3 && mode == "area"
+            adaptive_avg_pool1d(input, output_size)
+          elsif input.dim == 4 && mode == "area"
+            adaptive_avg_pool2d(input, output_size)
+          elsif input.dim == 5 && mode == "area"
+            adaptive_avg_pool3d(input, output_size)
+          elsif input.dim == 3 && mode == "linear"
+            # assert align_corners is not None
+            NN.upsample_linear1d(input, output_size, align_corners, sfl[0])
+          elsif input.dim == 3 && mode == "bilinear"
+            raise ArgumentError, "Got 3D input, but bilinear mode needs 4D input"
+          elsif input.dim == 3 && mode == "trilinear"
+            raise ArgumentError, "Got 3D input, but trilinear mode needs 5D input"
+          elsif input.dim == 4 && mode == "linear"
+            raise ArgumentError, "Got 4D input, but linear mode needs 3D input"
+          elsif input.dim == 4 && mode == "bilinear"
+            # assert align_corners is not None
+            NN.upsample_bilinear2d(input, output_size, align_corners, sfl[0], sfl[1])
+          elsif input.dim == 4 && mode == "trilinear"
+            raise ArgumentError, "Got 4D input, but trilinear mode needs 5D input"
+          elsif input.dim == 5 && mode == "linear"
+            raise ArgumentError, "Got 5D input, but linear mode needs 3D input"
+          elsif input.dim == 5 && mode == "bilinear"
+            raise ArgumentError, "Got 5D input, but bilinear mode needs 4D input"
+          elsif input.dim == 5 && mode == "trilinear"
+            # assert align_corners is not None
+            NN.upsample_trilinear3d(input, output_size, align_corners, sfl[0], sfl[1], sfl[2])
+          elsif input.dim == 4 && mode == "bicubic"
+            # assert align_corners is not None
+            NN.upsample_bicubic2d(input, output_size, align_corners, sfl[0], sfl[1])
+          else
+            raise ArgumentError, "Input Error: Only 3D, 4D and 5D input Tensors supported (got #{input.dim}D) for the modes: nearest | linear | bilinear | bicubic | trilinear (got #{mode})"
+          end
+        end
+
         private
 
         def softmax_dim(ndim)
@@ -480,6 +555,41 @@ module Torch
             out_size.zip(defaults.last(out_size.length)).map { |v, d| v || d }
           end
         end
+
+        def _interp_output_size(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"
+          end
+          if !size.nil? && !scale_factor.nil?
+            raise ArgumentError, "only one of size or scale_factor should be defined"
+          end
+          if !scale_factor.nil?
+            if scale_factor.is_a?(Array)
+              if scale_factor.length != dim
+                raise ArgumentError, "scale_factor shape must match input shape. Input is #{dim}D, scale_factor size is #{scale_factor.length}"
+              end
+            end
+          end
+
+          if !size.nil?
+            if size.is_a?(Array)
+              return size
+            else
+              return [size] * dim
+            end
+          end
+
+          raise "Failed assertion" if scale_factor.nil?
+          if scale_factor.is_a?(Array)
+            scale_factors = scale_factor
+          else
+            scale_factors = [scale_factor] * dim
+          end
+
+          dim.times.map { |i| (input.size(i + 2) * scale_factors[i]).floor }
+        end
       end
     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,15 @@ module Torch
             end
           end
         end
-        # TODO apply to more objects
+
+        @buffers.each_key do |k|
+          buf = @buffers[k]
+          unless buf.nil?
+            @buffers[k] = fn.call(buf)
+            instance_variable_set("@#{k}", @buffers[k])
+          end
+        end
+
         self
       end
 

data/lib/torch/nn/upsample.rb

@@ -0,0 +1,31 @@
+module Torch
+  module NN
+    class Upsample < Module
+      def initialize(size: nil, scale_factor: nil, mode: "nearest", align_corners: nil)
+        super()
+        @size = size
+        if scale_factor.is_a?(Array)
+          @scale_factor = scale_factor.map(&:to_f)
+        else
+          @scale_factor = scale_factor ? scale_factor.to_f : nil
+        end
+        @mode = mode
+        @align_corners = align_corners
+      end
+
+      def forward(input)
+        F.interpolate(input, size: @size, scale_factor: @scale_factor, mode: @mode, align_corners: @align_corners)
+      end
+
+      def extra_inspect
+        if !@scale_factor.nil?
+          info = "scale_factor: #{@scale_factor.inspect}"
+        else
+          info = "size: #{@size.inspect}"
+        end
+        info += ", mode: #{@mode.inspect}"
+        info
+      end
+    end
+  end
+end

data/lib/torch/tensor.rb

@@ -48,6 +48,11 @@ module Torch
     end
 
     def to(device = nil, dtype: nil, non_blocking: false, copy: false)
+      if device.is_a?(Symbol) && !dtype
+        dtype = device
+        device = nil
+      end
+
       device ||= self.device
       device = Device.new(device) if device.is_a?(String)
 
@@ -74,10 +79,6 @@ module Torch
       end
     end
 
-    def shape
-      dim.times.map { |i| size(i) }
-    end
-
     # mirror Python len()
     def length
       size(0)
@@ -103,11 +104,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]
@@ -124,9 +120,14 @@ module Torch
     end
 
     def type(dtype)
-      enum = DTYPE_TO_ENUM[dtype]
-      raise Error, "Unknown type: #{dtype}" unless enum
-      _type(enum)
+      if dtype.is_a?(Class)
+        raise Error, "Invalid type: #{dtype}" unless TENSOR_TYPE_CLASSES.include?(dtype)
+        dtype.new(self)
+      else
+        enum = DTYPE_TO_ENUM[dtype]
+        raise Error, "Invalid type: #{dtype}" unless enum
+        _type(enum)
+      end
     end
 
     def reshape(*size)
@@ -188,49 +189,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 +211,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 +227,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