torch-rb 0.1.4 → 0.1.5

data/ext/torch/torch_functions.hpp

@@ -0,0 +1,6 @@
+// generated by rake generate:functions
+// do not edit by hand
+
+#pragma once
+
+void add_torch_functions(Module m);

data/lib/torch.rb

@@ -1,6 +1,11 @@
 # ext
 require "torch/ext"
 
+# native functions
+require "torch/native/generator"
+require "torch/native/parser"
+require "torch/native/dispatcher"
+
 # modules
 require "torch/inspector"
 require "torch/tensor"
@@ -39,6 +44,10 @@ require "torch/nn/max_pool2d"
 require "torch/nn/avg_poolnd"
 require "torch/nn/avg_pool2d"
 
+# nn recurrent layers
+require "torch/nn/rnn_base"
+require "torch/nn/rnn"
+
 # nn linear layers
 require "torch/nn/bilinear"
 require "torch/nn/identity"
@@ -77,23 +86,23 @@ require "torch/nn/pairwise_distance"
 require "torch/nn/loss"
 require "torch/nn/weighted_loss"
 require "torch/nn/bce_loss"
-# require "torch/nn/bce_with_logits_loss"
-# require "torch/nn/cosine_embedding_loss"
+require "torch/nn/bce_with_logits_loss"
+require "torch/nn/cosine_embedding_loss"
 require "torch/nn/cross_entropy_loss"
 require "torch/nn/ctc_loss"
-# require "torch/nn/hinge_embedding_loss"
+require "torch/nn/hinge_embedding_loss"
 require "torch/nn/kl_div_loss"
 require "torch/nn/l1_loss"
-# require "torch/nn/margin_ranking_loss"
+require "torch/nn/margin_ranking_loss"
 require "torch/nn/mse_loss"
-# require "torch/nn/multi_label_margin_loss"
-# require "torch/nn/multi_label_soft_margin_loss"
-# require "torch/nn/multi_margin_loss"
+require "torch/nn/multi_label_margin_loss"
+require "torch/nn/multi_label_soft_margin_loss"
+require "torch/nn/multi_margin_loss"
 require "torch/nn/nll_loss"
 require "torch/nn/poisson_nll_loss"
-# require "torch/nn/smooth_l1_loss"
-# require "torch/nn/soft_margin_loss"
-# require "torch/nn/triplet_margin_loss"
+require "torch/nn/smooth_l1_loss"
+require "torch/nn/soft_margin_loss"
+require "torch/nn/triplet_margin_loss"
 
 # nn other
 require "torch/nn/functional"
@@ -142,6 +151,41 @@ module Torch
   }
   ENUM_TO_DTYPE = DTYPE_TO_ENUM.map(&:reverse).to_h
 
+  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?(Array)
+        Torch.tensor(args.first, dtype: dtype, device: device)
+      else
+        Torch.empty(*args, dtype: dtype, device: device)
+      end
+    end
+    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)
+
   class << self
     # Torch.float, Torch.long, etc
     DTYPE_TO_ENUM.each_key do |dtype|
@@ -191,6 +235,20 @@ module Torch
       }
     end
 
+    def no_grad
+      previous_value = grad_enabled?
+      begin
+        _set_grad_enabled(false)
+        yield
+      ensure
+        _set_grad_enabled(previous_value)
+      end
+    end
+
+    def device(str)
+      Device.new(str)
+    end
+
     # --- begin tensor creation: https://pytorch.org/cppdocs/notes/tensor_creation.html ---
 
     def arange(start, finish = nil, step = 1, **options)
@@ -308,26 +366,6 @@ module Torch
 
     # --- begin operations ---
 
-    %w(add sub mul div remainder).each do |op|
-      define_method(op) do |input, other, **options|
-        execute_op(op, input, other, **options)
-      end
-    end
-
-    def neg(input)
-      _neg(input)
-    end
-
-    def no_grad
-      previous_value = grad_enabled?
-      begin
-        _set_grad_enabled(false)
-        yield
-      ensure
-        _set_grad_enabled(previous_value)
-      end
-    end
-
     # TODO support out
     def mean(input, dim = nil, keepdim: false)
       if dim
@@ -346,34 +384,10 @@ module Torch
       end
     end
 
-    def argmax(input, dim = nil, keepdim: false)
-      if dim
-        _argmax_dim(input, dim, keepdim)
-      else
-        _argmax(input)
-      end
-    end
-
-    def eq(input, other)
-      _eq(input, other)
-    end
-
-    def norm(input)
-      _norm(input)
-    end
-
-    def pow(input, exponent)
-      _pow(input, exponent)
-    end
-
     def topk(input, k)
       _topk(input, k)
     end
 
-    def min(input)
-      _min(input)
-    end
-
     def max(input, dim = nil, keepdim: false, out: nil)
       if dim
         raise NotImplementedYet unless out
@@ -383,58 +397,6 @@ module Torch
       end
     end
 
-    def exp(input)
-      _exp(input)
-    end
-
-    def log(input)
-      _log(input)
-    end
-
-    def sign(input)
-      _sign(input)
-    end
-
-    def sigmoid(input)
-      _sigmoid(input)
-    end
-
-    def gt(input, other)
-      _gt(input, other)
-    end
-
-    def lt(input, other)
-      _lt(input, other)
-    end
-
-    def unsqueeze(input, dim)
-      _unsqueeze(input, dim)
-    end
-
-    def dot(input, tensor)
-      _dot(input, tensor)
-    end
-
-    def cat(tensors, dim = 0)
-      _cat(tensors, dim)
-    end
-
-    def matmul(input, other)
-      _matmul(input, other)
-    end
-
-    def reshape(input, shape)
-      _reshape(input, shape)
-    end
-
-    def flatten(input, start_dim: 0, end_dim: -1)
-      _flatten(input, start_dim, end_dim)
-    end
-
-    def sqrt(input)
-      _sqrt(input)
-    end
-
     # TODO make dim keyword argument
     def log_softmax(input, dim)
       _log_softmax(input, dim)
@@ -444,31 +406,8 @@ module Torch
       _softmax(input, dim)
     end
 
-    def abs(input)
-      _abs(input)
-    end
-
-    def device(str)
-      Device.new(str)
-    end
-
     private
 
-    def execute_op(op, input, other, out: nil)
-      scalar = other.is_a?(Numeric)
-      if out
-        # TODO make work with scalars
-        raise Error, "out not supported with scalar yet" if scalar
-        send("_#{op}_out", out, input, other)
-      else
-        if scalar
-          send("_#{op}_scalar", input, other)
-        else
-          send("_#{op}", input, other)
-        end
-      end
-    end
-
     def tensor_size(size)
       size.flatten
     end

data/lib/torch/native/dispatcher.rb

@@ -0,0 +1,48 @@
+# We use a generic interface for methods (*args, **options)
+# and this class to determine the C++ method to call
+#
+# This is needed since LibTorch uses function overloading,
+# which isn't available in Ruby or Python
+#
+# PyTorch uses this approach, but the parser/dispatcher is written in C++
+#
+# We could generate Ruby methods directly, but an advantage of this approach is
+# arguments and keyword arguments can be used interchangably like in Python,
+# making it easier to port code
+
+module Torch
+  module Native
+    module Dispatcher
+      class << self
+        def bind
+          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
+        end
+
+        def bind_functions(context, def_method, functions)
+          functions.group_by(&:ruby_name).sort_by { |g, _| g }.each do |name, funcs|
+            if def_method == :define_method
+              funcs.map! { |f| Function.new(f.function) }
+              funcs.each { |f| f.args.reject! { |a| a[:name] == "self" } }
+            end
+
+            defined = def_method == :define_method ? context.method_defined?(name) : context.respond_to?(name)
+            next if defined && name != "clone"
+
+            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])
+            end
+          end
+        end
+      end
+    end
+  end
+end
+
+Torch::Native::Dispatcher.bind

data/lib/torch/native/function.rb

@@ -0,0 +1,78 @@
+module Torch
+  module Native
+    class Function
+      attr_reader :function
+
+      def initialize(function)
+        @function = function
+      end
+
+      def func
+        @func ||= @function["func"]
+      end
+
+      def name
+        @name ||= func.split("(", 2).first
+      end
+
+      def python_module
+        @python_module ||= @function["python_module"]
+      end
+
+      def variants
+        @variants ||= (@function["variants"] || "function").split(", ")
+      end
+
+      def args
+        @args ||= begin
+          args = []
+          pos = true
+          args_str = func.split("(", 2).last.split(") ->").first
+          args_str.split(", ").each do |a|
+            if a == "*"
+              pos = false
+              next
+            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}
+          end
+          args
+        end
+      end
+
+      def out_size
+        @out_size ||= func.split("->").last.count("!")
+      end
+
+      def out?
+        out_size > 0 && base_name[-1] != "_"
+      end
+
+      def ruby_name
+        @ruby_name ||= begin
+          name = base_name
+          if name.end_with?("_")
+            "#{name[0..-2]}!"
+          elsif name.start_with?("is_")
+            "#{name[3..-1]}?"
+          else
+            name
+          end
+        end
+      end
+
+      def cpp_name
+        @cpp_name ||= "_" + name.downcase.sub(".", "_")
+      end
+
+      def base_name
+        @base_name ||= name.split(".").first
+      end
+    end
+  end
+end

data/lib/torch/native/generator.rb

@@ -0,0 +1,149 @@
+require "yaml"
+# use require_relative for
+# rake generate:function (without bundle)
+require_relative "function"
+
+module Torch
+  module Native
+    module Generator
+      class << self
+        def generate_cpp_functions
+          functions = grouped_functions
+          generate_cpp_file("torch", :define_singleton_method, functions[:torch])
+          generate_cpp_file("tensor", :define_method, functions[:tensor])
+          generate_cpp_file("nn", :define_singleton_method, functions[:nn])
+        end
+
+        def grouped_functions
+          functions = functions()
+
+          # remove 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) }
+          todo_functions, functions =
+            functions.partition do |f|
+              f.args.any? do |a|
+                a[:type].include?("?") && !["Tensor?", "Generator?", "int?"].include?(a[:type]) ||
+                skip_args.any? { |sa| a[:type].include?(sa) }
+              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.start_with?("avg_pool") || f.ruby_name == "cross"
+            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
+          # end
+
+          nn_functions, other_functions = functions.partition { |f| f.python_module == "nn" }
+          torch_functions = other_functions.select { |f| f.variants.include?("function") }
+          tensor_functions = other_functions.select { |f| f.variants.include?("method") }
+
+          {torch: torch_functions, tensor: tensor_functions, nn: nn_functions}
+        end
+
+        private
+
+        def generate_cpp_file(type, def_method, functions)
+          hpp_template = <<-TEMPLATE
+// generated by rake generate:functions
+// do not edit by hand
+
+#pragma once
+
+void add_%{type}_functions(Module m);
+        TEMPLATE
+
+          cpp_template = <<-TEMPLATE
+// generated by rake generate:functions
+// do not edit by hand
+
+#include <torch/torch.h>
+#include <rice/Module.hpp>
+#include "templates.hpp"
+
+void add_%{type}_functions(Module m) {
+  m
+  %{functions};
+}
+        TEMPLATE
+
+          cpp_defs = []
+          functions.sort_by(&:cpp_name).each do |func|
+            fargs = func.args.select { |a| a[:type] != "Generator?" }
+
+            cpp_args = []
+            fargs.each do |a|
+              t =
+                case a[:type]
+                when "Tensor"
+                  "const Tensor &"
+                when "Tensor?"
+                  # TODO better signature
+                  "OptionalTensor"
+                when "Tensor[]"
+                  "TensorList"
+                when "int"
+                  "int64_t"
+                when "float"
+                  "double"
+                when /\Aint\[/
+                  "IntArrayRef"
+                when /Tensor\(\S!?\)/
+                  "Tensor &"
+                else
+                  a[:type]
+                end
+
+              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
+
+            prefix = def_method == :define_method ? "self." : "torch::"
+
+            cpp_defs << ".#{def_method}(
+    \"#{func.cpp_name}\",
+    *[](#{cpp_args.join(", ")}) {
+      return #{prefix}#{dispatch}(#{args.join(", ")});
+    })"
+          end
+
+          hpp_contents = hpp_template % {type: type}
+          cpp_contents = cpp_template % {type: type, functions: cpp_defs.join("\n  ")}
+
+          path = File.expand_path("../../../ext/torch", __dir__)
+          File.write("#{path}/#{type}_functions.hpp", hpp_contents)
+          File.write("#{path}/#{type}_functions.cpp", cpp_contents)
+        end
+
+        def functions
+          @native_functions ||= YAML.load_file(path).map { |f| Function.new(f) }
+        end
+
+        def path
+          File.expand_path("native_functions.yaml", __dir__)
+        end
+      end
+    end
+  end
+end