torch-rb 0.1.3 → 0.1.8

data/lib/torch/native/parser.rb

@@ -0,0 +1,134 @@
+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[: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|
+            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
+
+        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] if values[fa[:name]].nil?
+          end
+
+          arg_types = func.args.map { |a| [a[:name], a[:type]] }.to_h
+
+          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.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
+              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/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

@@ -0,0 +1,19 @@
+module Torch
+  module NN
+    class AvgPool2d < 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_pool2d(input, @kernel_size, @stride, @padding, @ceil_mode, @count_include_pad, @divisor_override)
+      end
+    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

data/lib/torch/nn/avg_poolnd.rb

@@ -0,0 +1,9 @@
+module Torch
+  module NN
+    class AvgPoolNd < Module
+      def extra_inspect
+        format("kernel_size: %s, stride: %s, padding: %s", @kernel_size, @stride, @padding)
+      end
+    end
+  end
+end

data/lib/torch/nn/batch_norm.rb

@@ -0,0 +1,75 @@
+module Torch
+  module NN
+    class BatchNorm < Module
+      def initialize(num_features, eps: 1e-5, momentum: 0.1, affine: true, track_running_stats: true)
+        super()
+        @num_features = num_features
+        @eps = eps
+        @momentum = momentum
+        @affine = affine
+        @track_running_stats = track_running_stats
+        if @affine
+          @weight = Parameter.new(Torch::Tensor.new(num_features))
+          @bias = Parameter.new(Torch::Tensor.new(num_features))
+        else
+          register_parameter("weight", nil)
+          register_parameter("bias", nil)
+        end
+        if track_running_stats
+          register_buffer("running_mean", Torch.zeros(num_features))
+          register_buffer("running_var", Torch.ones(num_features))
+          register_buffer("num_batches_tracked", Torch.tensor(0, dtype: :long))
+        else
+          register_parameter("running_mean", nil)
+          register_parameter("running_var", nil)
+          register_parameter("num_batches_tracked", nil)
+        end
+        reset_parameters
+      end
+
+      def reset_running_stats
+        if @track_running_stats
+          @running_mean.zero!
+          @running_var.fill!(1)
+          @num_batches_tracked.zero!
+        end
+      end
+
+      def reset_parameters
+        reset_running_stats
+        if @affine
+          Init.ones!(@weight)
+          Init.zeros!(@bias)
+        end
+      end
+
+      def forward(input)
+        _check_input_dim(input)
+
+        if @momentum.nil?
+          exponential_average_factor = 0.0
+        else
+          exponential_average_factor = @momentum
+        end
+
+        if @training and @track_running_stats
+          if @num_batches_tracked.nil?
+            @num_batches_tracked += 1
+            if @momentum.nil?
+              exponential_average_factor = 1.0 / @num_batches_tracked.to_f
+            else
+              exponential_average_factor = @momentum
+            end
+          end
+        end
+
+        F.batch_norm(
+          input, @running_mean, @running_var,
+          weight: @weight, bias: @bias,
+          training: @training || !@track_running_stats,
+          momentum: exponential_average_factor, eps: @eps
+        )
+      end
+    end
+  end
+end

data/lib/torch/nn/batch_norm1d.rb

@@ -0,0 +1,11 @@
+module Torch
+  module NN
+    class BatchNorm1d < BatchNorm
+      def _check_input_dim(input)
+        if input.dim != 2 && input.dim != 3
+          raise ArgumentError, "expected 2D or 3D input (got #{input.dim}D input)"
+        end
+      end
+    end
+  end
+end

data/lib/torch/nn/batch_norm2d.rb

@@ -0,0 +1,11 @@
+module Torch
+  module NN
+    class BatchNorm2d < BatchNorm
+      def _check_input_dim(input)
+        if input.dim != 4
+          raise ArgumentError, "expected 4D input (got #{input.dim}D input)"
+        end
+      end
+    end
+  end
+end

data/lib/torch/nn/batch_norm3d.rb

@@ -0,0 +1,11 @@
+module Torch
+  module NN
+    class BatchNorm3d < BatchNorm
+      def _check_input_dim(input)
+        if input.dim != 5
+          raise ArgumentError, "expected 5D input (got #{input.dim}D input)"
+        end
+      end
+    end
+  end
+end

data/lib/torch/nn/bce_loss.rb

@@ -0,0 +1,13 @@
+module Torch
+  module NN
+    class BCELoss < WeightedLoss
+      def initialize(weight: nil, reduction: "mean")
+        super(weight, reduction)
+      end
+
+      def forward(input, target)
+        F.binary_cross_entropy(input, target, weight: @weight, reduction: @reduction)
+      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/bilinear.rb

@@ -0,0 +1,38 @@
+module Torch
+  module NN
+    class Bilinear < Module
+      def initialize(in1_features, in2_features, out_features, bias: true)
+        super()
+
+        @in1_features = in1_features
+        @in2_features = in2_features
+        @out_features = out_features
+        @weight = Parameter.new(Tensor.new(out_features, in1_features, in2_features))
+
+        if bias
+          @bias = Parameter.new(Tensor.new(out_features))
+        else
+          raise NotImplementedYet
+        end
+
+        reset_parameters
+      end
+
+      def reset_parameters
+        bound = 1 / Math.sqrt(@weight.size(1))
+        Init.uniform!(@weight, a: -bound, b: bound)
+        if @bias
+          Init.uniform!(@bias, a: -bound, b: bound)
+        end
+      end
+
+      def forward(input1, input2)
+        F.bilinear(input1, input2, @weight, @bias)
+      end
+
+      def extra_inspect
+        format("in1_features: %s, in2_features: %s, out_features: %s, bias: %s", @in1_features, @in2_features, @out_features, !@bias.nil?)
+      end
+    end
+  end
+end

data/lib/torch/nn/constant_pad1d.rb

@@ -0,0 +1,10 @@
+module Torch
+  module NN
+    class ConstantPad1d < ConstantPadNd
+      def initialize(padding, value)
+        super(value)
+        @padding = _pair(padding)
+      end
+    end
+  end
+end

data/lib/torch/nn/constant_pad2d.rb

@@ -0,0 +1,10 @@
+module Torch
+  module NN
+    class ConstantPad2d < ConstantPadNd
+      def initialize(padding, value)
+        super(value)
+        @padding = _quadrupal(padding)
+      end
+    end
+  end
+end

data/lib/torch/nn/constant_pad3d.rb

@@ -0,0 +1,10 @@
+module Torch
+  module NN
+    class ConstantPad3d < ConstantPadNd
+      def initialize(padding, value)
+        super(value)
+        @padding = _ntuple(6, padding)
+      end
+    end
+  end
+end

data/lib/torch/nn/constant_padnd.rb

@@ -0,0 +1,18 @@
+module Torch
+  module NN
+    class ConstantPadNd < Module
+      def initialize(value)
+        super()
+        @value = value
+      end
+
+      def forward(input)
+        F.pad(input, @padding, mode: "constant", value: @value)
+      end
+
+      def extra_inspect
+        format("padding: %s, value: %s", @padding, @value)
+      end
+    end
+  end
+end

data/lib/torch/nn/conv1d.rb

@@ -0,0 +1,22 @@
+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")
+
+        kernel_size = _single(kernel_size)
+        stride = _single(stride)
+        padding = _single(padding)
+        dilation = _single(dilation)
+        super(in_channels, out_channels, kernel_size, stride, padding, dilation, false, _single(0), groups, bias, padding_mode)
+      end
+
+      def forward(input)
+        if @padding_mode == "circular"
+          raise NotImplementedError
+        end
+        F.conv1d(input, @weight, @bias, @stride, @padding, @dilation, @groups)
+      end
+    end
+  end
+end