torch-rb 0.1.4 → 0.1.5

data/lib/torch/native/parser.rb

@@ -0,0 +1,97 @@
+module Torch
+  module Native
+    class Parser
+      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
+        @max_args = @functions.map { |f| f.args.count { |a| a[:pos] } }.max
+      end
+
+      def parse(args, options)
+        candidates = @functions.dup
+
+        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
+
+        # 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 }
+          end
+          # TODO show all bad keywords at once like Ruby?
+          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
+        candidates.select! do |func|
+          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]
+          end
+
+          arg_types = func.args.map { |a| [a[:name], a[:type]] }.to_h
+
+          values.each do |k, v|
+            t = arg_types[k].split("(").first
+            good =
+              case t
+              when "Tensor"
+                v.is_a?(Tensor)
+              when "Tensor[]"
+                v.all? { |v2| v2.is_a?(Tensor) }
+              when "int"
+                v.is_a?(Integer)
+              when "int[]"
+                v.all? { |v2| v2.is_a?(Integer) }
+              when "Scalar"
+                v.is_a?(Numeric)
+              when "bool"
+                v == true || v == false
+              else
+                raise Error, "Unknown argument type: #{arg_types[k]}. Please report a bug with #{@name}"
+              end
+
+            if !good
+              if candidates.size == 1
+                k = "input" if k == "self"
+                return {error: "#{@name}(): argument '#{k}' must be #{t}"}
+              end
+              break
+            end
+          end
+
+          if good
+            final_values = values
+          end
+
+          good
+        end
+
+        if candidates.size != 1
+          raise Error, "This should never happen. Please report a bug with #{@name}."
+        end
+
+        func = candidates.first
+        {
+          name: func.cpp_name,
+          args: func.args.map { |a| final_values[a[:name]] }
+        }
+      end
+    end
+  end
+end

data/lib/torch/nn/bce_with_logits_loss.rb

@@ -0,0 +1,15 @@
+module Torch
+  module NN
+    class BCEWithLogitsLoss < Loss
+      def initialize(weight: nil, reduction: "mean", pos_weight: nil)
+        super(reduction)
+        register_buffer("weight", weight)
+        register_buffer("pos_weight", pos_weight)
+      end
+
+      def forward(input, target)
+        F.binary_cross_entropy_with_logits(input, target, weight: weight, pos_weight: pos_weight, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/conv2d.rb

@@ -1,8 +1,6 @@
 module Torch
   module NN
     class Conv2d < ConvNd
-      attr_reader :bias, :weight
-
       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)

data/lib/torch/nn/cosine_embedding_loss.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class CosineEmbeddingLoss < Loss
+      def initialize(margin: 0, reduction: "mean")
+        super(reduction)
+        @margin = margin
+      end
+
+      def forward(input1, input2, target)
+        F.cosine_embedding_loss(input1, input2, target, margin: @margin, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/functional.rb

@@ -50,7 +50,8 @@ module Torch
           raise NotImplementedYet unless max_norm.nil? && norm_type == 2.0
 
           padding_idx ||= -1
-          Torch._embedding(input, weight, padding_idx, scale_grad_by_freq, sparse)
+          # weight and indices are swapped from Python interface
+          Torch._embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
         end
 
         def embedding_bag(input, weight, offsets: nil, max_norm: nil, norm_type: 2, scale_grad_by_freq: false, mode: "mean", sparse: false, per_sample_weights: nil)
@@ -76,8 +77,15 @@ module Torch
         # loss functions
 
         def binary_cross_entropy(input, target, weight: nil, reduction: "mean")
-          raise NotImplementedYet if weight
-          Torch.binary_cross_entropy(input, target, reduction)
+          NN._binary_cross_entropy(input, target, weight, reduction)
+        end
+
+        def binary_cross_entropy_with_logits(input, target, weight: nil, reduction: "mean", pos_weight: nil)
+          Torch._binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction)
+        end
+
+        def cosine_embedding_loss(input1, input2, target, margin: 0, reduction: "mean")
+          raise NotImplementedYet
         end
 
         def cross_entropy(input, target, weight: nil, ignore_index: -100, reduction: "mean")
@@ -86,28 +94,59 @@ module Torch
 
         def ctc_loss(log_probs, targets, input_lengths, target_lengths, blank: 0, reduction: "mean", zero_infinity: false)
           # call to_a on input_lengths and target_lengths for C++
-          Torch.ctc_loss(log_probs, targets, input_lengths.to_a, target_lengths.to_a, blank, reduction, zero_infinity)
+          Torch._ctc_loss_intlist(log_probs, targets, input_lengths.to_a, target_lengths.to_a, blank, reduction, zero_infinity)
+        end
+
+        def hinge_embedding_loss(input, target, margin: 1.0, reduction: "mean")
+          Torch._hinge_embedding_loss(input, target, margin, reduction)
         end
 
         def kl_div(input, target, reduction: "mean")
-          Torch.kl_div(input, target, reduction)
+          Torch._kl_div(input, target, reduction)
         end
 
         def l1_loss(input, target, reduction: "mean")
-          Torch.l1_loss(input, target, reduction)
+          NN._l1_loss(input, target, reduction)
+        end
+
+        def margin_ranking_loss(input1, input2, target, margin: 0, reduction: "mean")
+          raise NotImplementedYet
         end
 
         def mse_loss(input, target, reduction: "mean")
-          Torch.mse_loss(input, target, reduction)
+          NN._mse_loss(input, target, reduction)
+        end
+
+        def multilabel_margin_loss(input, target, reduction: "mean")
+          NN._multilabel_margin_loss(input, target, reduction)
+        end
+
+        def multilabel_soft_margin_loss(input, target, weight: nil)
+          raise NotImplementedYet
+        end
+
+        def multi_margin_loss(input, target, p: 1, margin: 1.0, weight: nil, reduction: "mean")
+          NN._multi_margin_loss(input, target, p, margin, weight, reduction)
         end
 
         def nll_loss(input, target, weight: nil, ignore_index: -100, reduction: "mean")
-          raise NotImplementedYet if weight
-          Torch.nll_loss(input, target, reduction, ignore_index)
+          NN._nll_loss(input, target, weight, reduction, ignore_index)
         end
 
         def poisson_nll_loss(input, target, log_input: true, full: false, eps: 1e-8, reduction: "mean")
-          Torch.poisson_nll_loss(input, target, log_input, full, eps, reduction)
+          Torch._poisson_nll_loss(input, target, log_input, full, eps, reduction)
+        end
+
+        def soft_margin_loss(input, target, reduction: "mean")
+          NN._soft_margin_loss(input, target, reduction)
+        end
+
+        def smooth_l1_loss(input, target, reduction: "mean")
+          NN._smooth_l1_loss(input, target, reduction)
+        end
+
+        def triplet_margin_loss(anchor, positive, negative, margin: 1.0, p: 2, eps: 1e-06, swap: false, reduction: "mean")
+          Torch._triplet_margin_loss(anchor, positive, negative, margin, p, eps, swap, reduction)
         end
 
         # end loss
@@ -123,7 +162,7 @@ module Torch
         end
 
         def softplus(input, beta: 1, threshold: 20)
-          Torch._softplus(input, beta, threshold)
+          NN._softplus(input, beta, threshold)
         end
 
         # TODO make dim keyword argument and update examples
@@ -134,7 +173,7 @@ module Torch
 
         def dropout(input, p: 0.5, training: true, inplace: false)
           if inplace
-            Torch._dropout!(input, p, training)
+            Torch._dropout_(input, p, training)
           else
             Torch._dropout(input, p, training)
           end
@@ -144,7 +183,7 @@ module Torch
           raise ArgumentError, "dropout probability has to be between 0 and 1, but got #{p}" if p < 0 || p > 1
 
           if inplace
-            Torch._feature_dropout!(input, p, training)
+            Torch._feature_dropout_(input, p, training)
           else
             Torch._feature_dropout(input, p, training)
           end
@@ -152,7 +191,7 @@ module Torch
 
         def dropout3d(input, p: 0.5, training: true, inplace: false)
           if inplace
-            Torch._feature_dropout!(input, p, training)
+            Torch._feature_dropout_(input, p, training)
           else
             Torch._feature_dropout(input, p, training)
           end
@@ -160,7 +199,7 @@ module Torch
 
         def alpha_dropout(input, p: 0.5, training: true, inplace: false)
           if inplace
-            Torch._alpha_dropout!(input, p, training)
+            Torch._alpha_dropout_(input, p, training)
           else
             Torch._alpha_dropout(input, p, training)
           end
@@ -168,7 +207,7 @@ module Torch
 
         def feature_alpha_dropout(input, p: 0.5, training: true, inplace: false)
           if inplace
-            Torch._feature_alpha_dropout!(input, p, training)
+            Torch._feature_alpha_dropout_(input, p, training)
           else
             Torch._feature_alpha_dropout(input, p, training)
           end

data/lib/torch/nn/hinge_embedding_loss.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class HingeEmbeddingLoss < Loss
+      def initialize(margin: 1.0, reduction: "mean")
+        super(reduction)
+        @margin = margin
+      end
+
+      def forward(input, target)
+        F.hinge_embedding_loss(input, target, margin: @margin, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/identity.rb

@@ -1,6 +1,7 @@
 module Torch
   module NN
     class Identity < Module
+      # written this way to support unused arguments
       def initialize(*args, **options)
         super()
       end

data/lib/torch/nn/margin_ranking_loss.rb

@@ -0,0 +1,14 @@
+module Torch
+  module NN
+    class MarginRankingLoss < Loss
+      def initialize(margin: 1.0, reduction: "mean")
+        super(reduction)
+        @margin = margin
+      end
+
+      def forward(input1, input2, target)
+        F.margin_ranking_loss(input1, input2, target, margin: @margin, reduction: @reduction)
+      end
+    end
+  end
+end

data/lib/torch/nn/module.rb

@@ -85,25 +85,56 @@ module Torch
       end
 
       def parameters
-        params = []
+        named_parameters.values
+      end
+
+      def named_parameters(prefix: "", recurse: true)
+        params = {}
+        if recurse
+          named_children.each do |name, mod|
+            params.merge!(mod.named_parameters(prefix: "#{name}.", recurse: recurse))
+          end
+        end
         instance_variables.each do |name|
           param = instance_variable_get(name)
-          params << param if param.is_a?(Parameter)
+          params[[prefix, name[1..-1]].join] = param if param.is_a?(Parameter)
+        end
+        @parameters.each do |name, param|
+          params[[prefix, name].join] = param
         end
-        params + modules.flat_map { |_, mod| mod.parameters }
+        params
+      end
+
+      def buffers
+        named_buffers.values
+      end
+
+      def named_buffers
+        @buffers || {}
       end
 
       def children
-        @modules.values
+        named_children.values
       end
 
-      def modules
+      def named_children
         modules = {}
         instance_variables.each do |name|
           mod = instance_variable_get(name)
           modules[name[1..-1]] = mod if mod.is_a?(Module)
         end
-        @modules.merge(modules)
+        @modules.each do |name, mod|
+          modules[name] = mod
+        end
+        modules
+      end
+
+      def modules
+        named_modules.values
+      end
+
+      def named_modules
+        {"" => self}.merge(named_children)
       end
 
       def train(mode = true)
@@ -140,12 +171,12 @@ module Torch
 
       def inspect
         name = self.class.name.split("::").last
-        if modules.empty?
+        if children.empty?
           "#{name}(#{extra_inspect})"
         else
           str = String.new
           str << "#{name}(\n"
-          modules.each do |name, mod|
+          children.each do |name, mod|
             str << "  (#{name}): #{mod.inspect}\n"
           end
           str << ")"
@@ -153,11 +184,21 @@ module Torch
       end
 
       def method_missing(method, *args, &block)
-        modules[method.to_s] || super
+        name = method.to_s
+        if named_parameters.key?(name)
+          named_parameters[name]
+        elsif named_buffers.key?(name)
+          named_buffers[name]
+        elsif named_modules.key?(name)
+          named_modules[name]
+        else
+          super
+        end
       end
 
       def respond_to?(method, include_private = false)
-        modules.key?(method.to_s) || super
+        name = method.to_s
+        named_parameters.key?(name) || named_buffers.key?(name) || named_modules.key?(name) || super
       end
 
       private
@@ -166,8 +207,14 @@ module Torch
         nil
       end
 
-      def format(str, *vars)
-        str % vars.map(&:inspect)
+      def format(str, *vars, **options)
+        vars =
+          if vars.any?
+            vars.map(&:inspect)
+          else
+            options.map { |k, v| [k, v.inspect] }.to_h
+          end
+        str % vars
       end
     end
   end

data/lib/torch/nn/multi_label_margin_loss.rb

@@ -0,0 +1,13 @@
+module Torch
+  module NN
+    class MultiLabelMarginLoss < Loss
+      def initialize(reduction: "mean")
+        super(reduction)
+      end
+
+      def forward(input, target)
+        F.multilabel_margin_loss(input, target, reduction: @reduction)
+      end
+    end
+  end
+end