torch-rb 0.3.2 → 0.3.7

data/lib/torch.rb

@@ -174,6 +174,9 @@ require "torch/nn/smooth_l1_loss"
 require "torch/nn/soft_margin_loss"
 require "torch/nn/triplet_margin_loss"
 
+# nn vision
+require "torch/nn/upsample"
+
 # nn other
 require "torch/nn/functional"
 require "torch/nn/init"
@@ -196,6 +199,32 @@ module Torch
     end
   end
 
+  # legacy
+  # but may make it easier to port tutorials
+  module Autograd
+    class Variable
+      def self.new(x)
+        raise ArgumentError, "Variable data has to be a tensor, but got #{x.class.name}" unless x.is_a?(Tensor)
+        warn "[torch] The Variable API is deprecated. Use tensors with requires_grad: true instead."
+        x
+      end
+    end
+  end
+
+  # TODO move to C++
+  class ByteStorage
+    # private
+    attr_reader :bytes
+
+    def initialize(bytes)
+      @bytes = bytes
+    end
+
+    def self.from_buffer(bytes)
+      new(bytes)
+    end
+  end
+
   # keys: https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.dtype
   # values: https://github.com/pytorch/pytorch/blob/master/c10/core/ScalarType.h
   DTYPE_TO_ENUM = {
@@ -224,40 +253,43 @@ module Torch
   }
   ENUM_TO_DTYPE = DTYPE_TO_ENUM.map(&:reverse).to_h
 
+  TENSOR_TYPE_CLASSES = []
+
   def self._make_tensor_class(dtype, cuda = false)
     cls = Class.new
     device = cuda ? "cuda" : "cpu"
     cls.define_singleton_method("new") do |*args|
       if args.size == 1 && args.first.is_a?(Tensor)
         args.first.send(dtype).to(device)
+      elsif args.size == 1 && args.first.is_a?(ByteStorage) && dtype == :uint8
+        bytes = args.first.bytes
+        Torch._from_blob(bytes, [bytes.bytesize], TensorOptions.new.dtype(DTYPE_TO_ENUM[dtype]))
       elsif args.size == 1 && args.first.is_a?(Array)
         Torch.tensor(args.first, dtype: dtype, device: device)
       else
         Torch.empty(*args, dtype: dtype, device: device)
       end
     end
+    TENSOR_TYPE_CLASSES << cls
     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
@@ -388,6 +420,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
 
@@ -460,6 +496,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/native/dispatcher.rb

@@ -22,21 +22,43 @@ module Torch
         end
 
         def bind_functions(context, def_method, functions)
+          instance_method = def_method == :define_method
           functions.group_by(&:ruby_name).sort_by { |g, _| g }.each do |name, funcs|
-            if def_method == :define_method
+            if instance_method
               funcs.map! { |f| Function.new(f.function) }
-              funcs.each { |f| f.args.reject! { |a| a[:name] == "self" } }
+              funcs.each { |f| f.args.reject! { |a| a[:name] == :self } }
             end
 
-            defined = def_method == :define_method ? context.method_defined?(name) : context.respond_to?(name)
+            defined = instance_method ? context.method_defined?(name) : context.respond_to?(name)
             next if defined && name != "clone"
 
-            parser = Parser.new(funcs)
+            # skip parser when possible for performance
+            if funcs.size == 1 && funcs.first.args.size == 0
+              # functions with no arguments
+              if instance_method
+                context.send(:alias_method, name, funcs.first.cpp_name)
+              else
+                context.singleton_class.send(:alias_method, name, funcs.first.cpp_name)
+              end
+            elsif funcs.size == 2 && funcs.map { |f| f.arg_types.values }.sort == [["Scalar"], ["Tensor"]]
+              # functions that take a tensor or scalar
+              scalar_name, tensor_name = funcs.sort_by { |f| f.arg_types.values }.map(&:cpp_name)
+              context.send(def_method, name) do |other|
+                case other
+                when Tensor
+                  send(tensor_name, other)
+                else
+                  send(scalar_name, other)
+                end
+              end
+            else
+              parser = Parser.new(funcs)
 
-            context.send(def_method, name) do |*args, **options|
-              result = parser.parse(args, options)
-              raise ArgumentError, result[:error] if result[:error]
-              send(result[:name], *result[:args])
+              context.send(def_method, name) do |*args, **options|
+                result = parser.parse(args, options)
+                raise ArgumentError, result[:error] if result[:error]
+                send(result[:name], *result[:args])
+              end
             end
           end
         end

data/lib/torch/native/function.rb

@@ -1,10 +1,15 @@
 module Torch
   module Native
     class Function
-      attr_reader :function
+      attr_reader :function, :tensor_options
 
       def initialize(function)
         @function = function
+
+        # note: don't modify function in-place
+        @tensor_options_str = ", *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None)"
+        @tensor_options = @function["func"].include?(@tensor_options_str)
+        @out = out_size > 0 && base_name[-1] != "_"
       end
 
       def func
@@ -27,7 +32,7 @@ module Torch
         @args ||= begin
           args = []
           pos = true
-          args_str = func.split("(", 2).last.split(") ->").first
+          args_str = func.sub(@tensor_options_str, ")").split("(", 2).last.split(") ->").first
           args_str.split(", ").each do |a|
             if a == "*"
               pos = false
@@ -68,12 +73,88 @@ module Torch
             next if t == "Generator?"
             next if t == "MemoryFormat"
             next if t == "MemoryFormat?"
-            args << {name: k, type: t, default: d, pos: pos, has_default: has_default}
+            args << {name: k.to_sym, type: t, default: d, pos: pos, has_default: has_default}
           end
           args
         end
       end
 
+      def arg_checkers
+        @arg_checkers ||= begin
+          checkers = {}
+          arg_types.each do |k, t|
+            checker =
+              case t
+              when "Tensor"
+                ->(v) { v.is_a?(Tensor) }
+              when "Tensor?"
+                ->(v) { v.nil? || v.is_a?(Tensor) }
+              when "Tensor[]", "Tensor?[]"
+                ->(v) { v.is_a?(Array) && v.all? { |v2| v2.is_a?(Tensor) } }
+              when "int"
+                if k == :reduction
+                  ->(v) { v.is_a?(String) }
+                else
+                  ->(v) { v.is_a?(Integer) }
+                end
+              when "int?"
+                ->(v) { v.is_a?(Integer) || v.nil? }
+              when "float?"
+                ->(v) { v.is_a?(Numeric) || v.nil? }
+              when "bool?"
+                ->(v) { v == true || v == false || v.nil? }
+              when "float"
+                ->(v) { v.is_a?(Numeric) }
+              when /int\[.*\]/
+                ->(v) { v.is_a?(Array) && v.all? { |v2| v2.is_a?(Integer) } }
+              when "Scalar"
+                ->(v) { v.is_a?(Numeric) }
+              when "Scalar?"
+                ->(v) { v.is_a?(Numeric) || v.nil? }
+              when "ScalarType"
+                ->(v) { false } # not supported yet
+              when "ScalarType?"
+                ->(v) { v.nil? }
+              when "bool"
+                ->(v) { v == true || v == false }
+              when "str"
+                ->(v) { v.is_a?(String) }
+              else
+                raise Error, "Unknown argument type: #{t}. Please report a bug with #{@name}."
+              end
+            checkers[k] = checker
+          end
+          checkers
+        end
+      end
+
+      def int_array_lengths
+        @int_array_lengths ||= begin
+          ret = {}
+          arg_types.each do |k, t|
+            if t.match?(/\Aint\[.+\]\z/)
+              size = t[4..-2]
+              raise Error, "Unknown size: #{size}. Please report a bug with #{@name}." unless size =~ /\A\d+\z/
+              ret[k] = size.to_i
+            end
+          end
+          ret
+        end
+      end
+
+      def arg_names
+        @arg_names ||= args.map { |a| a[:name] }
+      end
+
+      def arg_types
+        @arg_types ||= args.map { |a| [a[:name], a[:type].split("(").first] }.to_h
+      end
+
+      def arg_defaults
+        # TODO find out why can't use select here
+        @arg_defaults ||= args.map { |a| [a[:name], a[:default]] }.to_h
+      end
+
       def out_size
         @out_size ||= func.split("->").last.count("!")
       end
@@ -82,8 +163,16 @@ module Torch
         @ret_size ||= func.split("->").last.split(", ").size
       end
 
+      def ret_array?
+        @ret_array ||= func.split("->").last.include?('[]')
+      end
+
+      def ret_void?
+        func.split("->").last.strip == "()"
+      end
+
       def out?
-        out_size > 0 && base_name[-1] != "_"
+        @out
       end
 
       def ruby_name

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
@@ -89,15 +72,18 @@ void add_%{type}_functions(Module m);
 #include <rice/Module.hpp>
 #include "templates.hpp"
 
+%{functions}
+
 void add_%{type}_functions(Module m) {
-  m
-  %{functions};
+  %{add_functions}
 }
         TEMPLATE
 
           cpp_defs = []
+          add_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,52 +95,70 @@ void add_%{type}_functions(Module m) {
                   # TODO better signature
                   "OptionalTensor"
                 when "ScalarType?"
-                  "OptionalScalarType"
-                when "Tensor[]"
-                  "TensorList"
-                when "Tensor?[]"
+                  "torch::optional<ScalarType>"
+                when "Tensor[]", "Tensor?[]"
                   # TODO make optional
-                  "TensorList"
+                  "std::vector<Tensor>"
                 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\[/
-                  "IntArrayRef"
+                  "std::vector<int64_t>"
                 when /Tensor\(\S!?\)/
                   "Tensor &"
                 when "str"
                   "std::string"
-                else
+                when "TensorOptions"
+                  "const torch::TensorOptions &"
+                when "Layout?"
+                  "torch::optional<Layout>"
+                when "Device?"
+                  "torch::optional<Device>"
+                when "Scalar", "bool", "ScalarType", "Layout", "Device", "Storage"
                   a[:type]
+                else
+                  raise "Unknown type: #{a[:type]}"
                 end
 
-              t = "MyReduction" if a[:name] == "reduction" && t == "int64_t"
+              t = "MyReduction" if a[:name] == :reduction && t == "int64_t"
               cpp_args << [t, a[:name]].join(" ").sub("& ", "&")
             end
 
             dispatch = func.out? ? "#{func.base_name}_out" : func.base_name
             args = fargs.map { |a| a[:name] }
             args.unshift(*args.pop(func.out_size)) if func.out?
-            args.delete("self") if def_method == :define_method
+            args.delete(:self) if def_method == :define_method
 
             prefix = def_method == :define_method ? "self." : "torch::"
 
             body = "#{prefix}#{dispatch}(#{args.join(", ")})"
-            # TODO check type as well
-            if func.ret_size > 1
+
+            if func.cpp_name == "_fill_diagonal_"
+              body = "to_ruby<torch::Tensor>(#{body})"
+            elsif !func.ret_void?
               body = "wrap(#{body})"
             end
 
-            cpp_defs << ".#{def_method}(
-    \"#{func.cpp_name}\",
-    *[](#{cpp_args.join(", ")}) {
-      return #{body};
-    })"
+            cpp_defs << "// #{func.func}
+static #{func.ret_void? ? "void" : "Object"} #{type}#{func.cpp_name}(#{cpp_args.join(", ")})
+{
+  return #{body};
+}"
+
+            add_defs << "m.#{def_method}(\"#{func.cpp_name}\", #{type}#{func.cpp_name});"
           end
 
           hpp_contents = hpp_template % {type: type}
-          cpp_contents = cpp_template % {type: type, functions: cpp_defs.join("\n  ")}
+          cpp_contents = cpp_template % {type: type, functions: cpp_defs.join("\n\n"), add_functions: add_defs.join("\n  ")}
 
           path = File.expand_path("../../../ext/torch", __dir__)
           File.write("#{path}/#{type}_functions.hpp", hpp_contents)