torch-rb 0.1.5 → 0.1.6

data/lib/torch.rb

@@ -29,6 +29,7 @@ require "torch/optim/lr_scheduler/step_lr"
 
 # nn parameters
 require "torch/nn/parameter"
+require "torch/nn/utils"
 
 # nn containers
 require "torch/nn/module"
@@ -36,17 +37,61 @@ require "torch/nn/sequential"
 
 # nn convolution layers
 require "torch/nn/convnd"
+require "torch/nn/conv1d"
 require "torch/nn/conv2d"
+require "torch/nn/conv3d"
+require "torch/nn/unfold"
+require "torch/nn/fold"
 
 # nn pooling layers
 require "torch/nn/max_poolnd"
+require "torch/nn/max_pool1d"
 require "torch/nn/max_pool2d"
+require "torch/nn/max_pool3d"
+require "torch/nn/max_unpoolnd"
+require "torch/nn/max_unpool1d"
+require "torch/nn/max_unpool2d"
+require "torch/nn/max_unpool3d"
 require "torch/nn/avg_poolnd"
+require "torch/nn/avg_pool1d"
 require "torch/nn/avg_pool2d"
+require "torch/nn/avg_pool3d"
+require "torch/nn/lp_poolnd"
+require "torch/nn/lp_pool1d"
+require "torch/nn/lp_pool2d"
+
+# nn padding layers
+require "torch/nn/reflection_padnd"
+require "torch/nn/reflection_pad1d"
+require "torch/nn/reflection_pad2d"
+require "torch/nn/replication_padnd"
+require "torch/nn/replication_pad1d"
+require "torch/nn/replication_pad2d"
+require "torch/nn/replication_pad3d"
+require "torch/nn/constant_padnd"
+require "torch/nn/constant_pad1d"
+require "torch/nn/constant_pad2d"
+require "torch/nn/constant_pad3d"
+require "torch/nn/zero_pad2d"
+
+# nn normalization layers
+require "torch/nn/batch_norm"
+require "torch/nn/batch_norm1d"
+require "torch/nn/batch_norm2d"
+require "torch/nn/batch_norm3d"
+require "torch/nn/group_norm"
+require "torch/nn/instance_norm"
+require "torch/nn/instance_norm1d"
+require "torch/nn/instance_norm2d"
+require "torch/nn/instance_norm3d"
+require "torch/nn/layer_norm"
+require "torch/nn/local_response_norm"
 
 # nn recurrent layers
 require "torch/nn/rnn_base"
 require "torch/nn/rnn"
+require "torch/nn/lstm"
+require "torch/nn/gru"
 
 # nn linear layers
 require "torch/nn/bilinear"
@@ -62,11 +107,17 @@ require "torch/nn/dropout3d"
 require "torch/nn/feature_alpha_dropout"
 
 # nn activations
+require "torch/nn/hardshrink"
 require "torch/nn/leaky_relu"
+require "torch/nn/log_sigmoid"
 require "torch/nn/prelu"
 require "torch/nn/relu"
 require "torch/nn/sigmoid"
 require "torch/nn/softplus"
+require "torch/nn/softshrink"
+require "torch/nn/softsign"
+require "torch/nn/tanh"
+require "torch/nn/tanhshrink"
 
 # nn activations other
 require "torch/nn/log_softmax"
@@ -364,48 +415,6 @@ module Torch
       zeros(input.size, like_options(input, options))
     end
 
-    # --- begin operations ---
-
-    # TODO support out
-    def mean(input, dim = nil, keepdim: false)
-      if dim
-        _mean_dim(input, dim, keepdim)
-      else
-        _mean(input)
-      end
-    end
-
-    # TODO support dtype
-    def sum(input, dim = nil, keepdim: false)
-      if dim
-        _sum_dim(input, dim, keepdim)
-      else
-        _sum(input)
-      end
-    end
-
-    def topk(input, k)
-      _topk(input, k)
-    end
-
-    def max(input, dim = nil, keepdim: false, out: nil)
-      if dim
-        raise NotImplementedYet unless out
-        _max_out(out[0], out[1], input, dim, keepdim)
-      else
-        _max(input)
-      end
-    end
-
-    # TODO make dim keyword argument
-    def log_softmax(input, dim)
-      _log_softmax(input, dim)
-    end
-
-    def softmax(input, dim: nil)
-      _softmax(input, dim)
-    end
-
     private
 
     def tensor_size(size)

data/lib/torch/native/dispatcher.rb

@@ -18,7 +18,7 @@ module Torch
           functions = Generator.grouped_functions
           bind_functions(::Torch, :define_singleton_method, functions[:torch])
           bind_functions(::Torch::Tensor, :define_method, functions[:tensor])
-          # NN functions are internal, so no need to bind
+          bind_functions(::Torch::NN, :define_singleton_method, functions[:nn])
         end
 
         def bind_functions(context, def_method, functions)

data/lib/torch/native/function.rb

@@ -35,11 +35,38 @@ module Torch
             end
             t, _, k = a.rpartition(" ")
             k, d = k.split("=")
-            d = d.to_i if d.to_i.to_s == d
-            d = true if d == "True"
-            d = false if d == "False"
-            d = nil if d == "None"
-            args << {name: k, type: t, default: d, pos: pos}
+            has_default = !d.nil?
+
+            if d
+              d =
+                case d
+                when "True"
+                  true
+                when "False"
+                  false
+                when "None"
+                  nil
+                when /\A\-?\d+\z/
+                  d.to_i
+                when "[]"
+                  []
+                when "[0,1]"
+                  [0, 1]
+                when /\A\de\-\d+\z/, /\A\d+\.\d+\z/
+                  d.to_f
+                when "Mean"
+                  "mean"
+                when "contiguous_format"
+                  d
+                when "long"
+                  :long
+                else
+                  raise "Unknown default: #{d}"
+                end
+            end
+
+            next if t == "Generator?"
+            args << {name: k, type: t, default: d, pos: pos, has_default: has_default}
           end
           args
         end
@@ -49,6 +76,10 @@ module Torch
         @out_size ||= func.split("->").last.count("!")
       end
 
+      def ret_size
+        @ret_size ||= func.split("->").last.split(", ").size
+      end
+
       def out?
         out_size > 0 && base_name[-1] != "_"
       end

data/lib/torch/native/generator.rb

@@ -17,14 +17,23 @@ module Torch
         def grouped_functions
           functions = functions()
 
-          # remove functions
+          # skip functions
           skip_binding = ["unique_dim_consecutive", "einsum", "normal"]
-          skip_args = ["bool[3]", "Dimname", "ScalarType", "MemoryFormat", "Storage", "ConstQuantizerPtr"]
-          functions.reject! { |f| f.ruby_name.start_with?("_") || f.ruby_name.end_with?("_backward") || skip_binding.include?(f.ruby_name) }
+          skip_args = ["bool[3]", "Dimname", "MemoryFormat", "Layout", "Storage", "ConstQuantizerPtr"]
+
+          # remove functions
+          functions.reject! do |f|
+            f.ruby_name.start_with?("_") ||
+            f.ruby_name.end_with?("_backward") ||
+            skip_binding.include?(f.ruby_name) ||
+            f.args.any? { |a| a[:type].include?("Dimname") }
+          end
+
+          # separate out into todo
           todo_functions, functions =
             functions.partition do |f|
               f.args.any? do |a|
-                a[:type].include?("?") && !["Tensor?", "Generator?", "int?"].include?(a[:type]) ||
+                a[:type].include?("?") && !["Tensor?", "Generator?", "int?", "ScalarType?"].include?(a[:type]) ||
                 skip_args.any? { |sa| a[:type].include?(sa) }
               end
             end
@@ -33,7 +42,9 @@ module Torch
           # 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.start_with?("avg_pool") || f.ruby_name == "cross"
+            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?(".") ? "_" : "."
@@ -85,7 +96,7 @@ 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 #.select { |a| a[:type] != "Generator?" }
 
             cpp_args = []
             fargs.each do |a|
@@ -96,6 +107,8 @@ void add_%{type}_functions(Module m) {
                 when "Tensor?"
                   # TODO better signature
                   "OptionalTensor"
+                when "ScalarType?"
+                  "OptionalScalarType"
                 when "Tensor[]"
                   "TensorList"
                 when "int"
@@ -121,10 +134,16 @@ 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
+              body = "wrap(#{body})"
+            end
+
             cpp_defs << ".#{def_method}(
     \"#{func.cpp_name}\",
     *[](#{cpp_args.join(", ")}) {
-      return #{prefix}#{dispatch}(#{args.join(", ")});
+      return #{body};
     })"
           end
 

data/lib/torch/native/parser.rb

@@ -4,19 +4,44 @@ module Torch
       def initialize(functions)
         @functions = functions
         @name = @functions.first.ruby_name
-        @min_args = @functions.map { |f| f.args.count { |a| a[:pos] && a[:default].nil? } }.min
+        @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
       end
 
       def parse(args, options)
         candidates = @functions.dup
 
+        # remove nil
+        while args.any? && args.last.nil?
+          args.pop
+        end
+
+        # TODO account for args passed as options here
         if args.size < @min_args || args.size > @max_args
           expected = String.new(@min_args.to_s)
           expected += "..#{@max_args}" if @max_args != @min_args
           return {error: "wrong number of arguments (given #{args.size}, expected #{expected})"}
         end
 
+        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|
@@ -26,12 +51,6 @@ module Torch
           return {error: "unknown keyword: #{k}"} if candidates.empty?
         end
 
-        # exclude functions missing required options
-        candidates.reject! do |func|
-          # TODO make more generic
-          func.out? && !options[:out]
-        end
-
         final_values = {}
 
         # check args
@@ -41,29 +60,47 @@ module Torch
           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]
+            values[fa[:name]] = fa[:default] if values[fa[:name]].nil?
           end
 
           arg_types = func.args.map { |a| [a[:name], a[:type]] }.to_h
 
-          values.each do |k, v|
+          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[]"
-                v.all? { |v2| v2.is_a?(Tensor) }
+                v.is_a?(Array) && v.all? { |v2| v2.is_a?(Tensor) }
               when "int"
-                v.is_a?(Integer)
-              when "int[]"
-                v.all? { |v2| v2.is_a?(Integer) }
+                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
               else
-                raise Error, "Unknown argument type: #{arg_types[k]}. Please report a bug with #{@name}"
+                raise Error, "Unknown argument type: #{arg_types[k]}. Please report a bug with #{@name}."
               end
 
             if !good

data/lib/torch/nn/avg_pool1d.rb

@@ -0,0 +1,18 @@
+module Torch
+  module NN
+    class AvgPool1d < AvgPoolNd
+      def initialize(kernel_size, stride: nil, padding: 0, ceil_mode: false, count_include_pad: true)
+        super()
+        @kernel_size = _single(kernel_size)
+        @stride = _single(stride || kernel_size)
+        @padding = _single(padding)
+        @ceil_mode = ceil_mode
+        @count_include_pad = count_include_pad
+      end
+
+      def forward(input)
+        F.avg_pool1d(input, @kernel_size, @stride, @padding, @ceil_mode, @count_include_pad)
+      end
+    end
+  end
+end

data/lib/torch/nn/avg_pool2d.rb

@@ -1,13 +1,18 @@
 module Torch
   module NN
     class AvgPool2d < AvgPoolNd
-      def initialize(kernel_size)
+      def initialize(kernel_size, stride: nil, padding: 0, ceil_mode: false, count_include_pad: true, divisor_override: nil)
         super()
         @kernel_size = kernel_size
+        @stride = stride || kernel_size
+        @padding = padding
+        @ceil_mode = ceil_mode
+        @count_include_pad = count_include_pad
+        @divisor_override = divisor_override
       end
 
       def forward(input)
-        F.avg_pool2d(input, @kernel_size)
+        F.avg_pool2d(input, @kernel_size, @stride, @padding, @ceil_mode, @count_include_pad, @divisor_override)
       end
     end
   end

data/lib/torch/nn/avg_pool3d.rb

@@ -0,0 +1,19 @@
+module Torch
+  module NN
+    class AvgPool3d < AvgPoolNd
+      def initialize(kernel_size, stride: nil, padding: 0, ceil_mode: false, count_include_pad: true, divisor_override: nil)
+        super()
+        @kernel_size = kernel_size
+        @stride = stride || kernel_size
+        @padding = padding
+        @ceil_mode = ceil_mode
+        @count_include_pad = count_include_pad
+        @divisor_override = divisor_override
+      end
+
+      def forward(input)
+        F.avg_pool3d(input, @kernel_size, @stride, @padding, @ceil_mode, @count_include_pad, @divisor_override)
+      end
+    end
+  end
+end