torch-rb 0.1.3 → 0.1.8

data/ext/torch/extconf.rb

@@ -10,6 +10,9 @@ $CXXFLAGS << " -std=c++11"
 # silence ruby/intern.h warning
 $CXXFLAGS << " -Wno-deprecated-register"
 
+# silence torch warnings
+$CXXFLAGS << " -Wno-shorten-64-to-32 -Wno-missing-noreturn"
+
 inc, lib = dir_config("torch")
 
 inc ||= "/usr/local/include"
@@ -22,4 +25,10 @@ $LDFLAGS << " -Wl,-rpath,#{lib}"
 $LDFLAGS << " -L#{lib}"
 $LDFLAGS << " -ltorch -lc10"
 
+# generate C++ functions
+puts "Generating C++ functions..."
+require_relative "../../lib/torch/native/generator"
+Torch::Native::Generator.generate_cpp_functions
+
+# create makefile
 create_makefile("torch/ext")

data/ext/torch/templates.cpp

@@ -0,0 +1,55 @@
+#include <torch/torch.h>
+#include <rice/Object.hpp>
+#include "templates.hpp"
+
+Object wrap(std::tuple<torch::Tensor, torch::Tensor> x) {
+  Array a;
+  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
+  return Object(a);
+}
+
+Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> x) {
+  Array a;
+  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<2>(x)));
+  return Object(a);
+}
+
+Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> x) {
+  Array a;
+  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<2>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<3>(x)));
+  return Object(a);
+}
+
+Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> x) {
+  Array a;
+  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<2>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<3>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<4>(x)));
+  return Object(a);
+}
+
+Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, int64_t> x) {
+  Array a;
+  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<2>(x)));
+  a.push(to_ruby<int64_t>(std::get<3>(x)));
+  return Object(a);
+}
+
+Object wrap(std::tuple<torch::Tensor, torch::Tensor, double, int64_t> x) {
+  Array a;
+  a.push(to_ruby<torch::Tensor>(std::get<0>(x)));
+  a.push(to_ruby<torch::Tensor>(std::get<1>(x)));
+  a.push(to_ruby<double>(std::get<2>(x)));
+  a.push(to_ruby<int64_t>(std::get<3>(x)));
+  return Object(a);
+}

data/ext/torch/templates.hpp

@@ -0,0 +1,244 @@
+#pragma once
+
+#ifdef isfinite
+#undef isfinite
+#endif
+
+#include <rice/Array.hpp>
+#include <rice/Object.hpp>
+
+using namespace Rice;
+
+// need to wrap torch::IntArrayRef() since
+// it doesn't own underlying data
+class IntArrayRef {
+  std::vector<int64_t> vec;
+  public:
+    IntArrayRef(Object o) {
+      Array a = Array(o);
+      for (size_t i = 0; i < a.size(); i++) {
+        vec.push_back(from_ruby<int64_t>(a[i]));
+      }
+    }
+    operator torch::IntArrayRef() {
+      return torch::IntArrayRef(vec);
+    }
+};
+
+template<>
+inline
+IntArrayRef from_ruby<IntArrayRef>(Object x)
+{
+  return IntArrayRef(x);
+}
+
+// for now
+class Scalar {
+  torch::Scalar value;
+  public:
+    Scalar(Object o) {
+      // TODO cast based on Ruby type
+      if (o.rb_type() == T_FIXNUM) {
+        value = torch::Scalar(from_ruby<int64_t>(o));
+      } else {
+        value = torch::Scalar(from_ruby<float>(o));
+      }
+    }
+    operator torch::Scalar() {
+      return value;
+    }
+};
+
+template<>
+inline
+Scalar from_ruby<Scalar>(Object x)
+{
+  return Scalar(x);
+}
+
+class TensorList {
+  std::vector<torch::Tensor> vec;
+  public:
+    TensorList(Object o) {
+      Array a = Array(o);
+      for (size_t i = 0; i < a.size(); i++) {
+        vec.push_back(from_ruby<torch::Tensor>(a[i]));
+      }
+    }
+    operator torch::TensorList() {
+      return torch::TensorList(vec);
+    }
+};
+
+template<>
+inline
+TensorList from_ruby<TensorList>(Object x)
+{
+  return TensorList(x);
+}
+
+class FanModeType {
+  std::string s;
+  public:
+    FanModeType(Object o) {
+      s = String(o).str();
+    }
+    operator torch::nn::init::FanModeType() {
+      if (s == "fan_in") {
+        return torch::kFanIn;
+      } else if (s == "fan_out") {
+        return torch::kFanOut;
+      } else {
+        throw std::runtime_error("Unsupported nonlinearity type: " + s);
+      }
+    }
+};
+
+template<>
+inline
+FanModeType from_ruby<FanModeType>(Object x)
+{
+  return FanModeType(x);
+}
+
+class NonlinearityType {
+  std::string s;
+  public:
+    NonlinearityType(Object o) {
+      s = String(o).str();
+    }
+    operator torch::nn::init::NonlinearityType() {
+      if (s == "linear") {
+        return torch::kLinear;
+      } else if (s == "conv1d") {
+        return torch::kConv1D;
+      } else if (s == "conv2d") {
+        return torch::kConv2D;
+      } else if (s == "conv3d") {
+        return torch::kConv3D;
+      } else if (s == "conv_transpose1d") {
+        return torch::kConvTranspose1D;
+      } else if (s == "conv_transpose2d") {
+        return torch::kConvTranspose2D;
+      } else if (s == "conv_transpose3d") {
+        return torch::kConvTranspose3D;
+      } else if (s == "sigmoid") {
+        return torch::kSigmoid;
+      } else if (s == "tanh") {
+        return torch::kTanh;
+      } else if (s == "relu") {
+        return torch::kReLU;
+      } else if (s == "leaky_relu") {
+        return torch::kLeakyReLU;
+      } else {
+        throw std::runtime_error("Unsupported nonlinearity type: " + s);
+      }
+    }
+};
+
+template<>
+inline
+NonlinearityType from_ruby<NonlinearityType>(Object x)
+{
+  return NonlinearityType(x);
+}
+
+class MyReduction {
+  Object value;
+  public:
+    MyReduction(Object o) {
+      value = o;
+    }
+    operator int64_t() {
+      if (value.is_nil()) {
+        return torch::Reduction::None;
+      }
+
+      std::string s = String(value).str();
+      if (s == "mean") {
+        return torch::Reduction::Mean;
+      } else if (s == "sum") {
+        return torch::Reduction::Sum;
+      } else {
+        throw std::runtime_error("Unsupported reduction: " + s);
+      }
+    }
+};
+
+template<>
+inline
+MyReduction from_ruby<MyReduction>(Object x)
+{
+  return MyReduction(x);
+}
+
+typedef torch::Tensor Tensor;
+
+class OptionalTensor {
+  Object value;
+  public:
+    OptionalTensor(Object o) {
+      value = o;
+    }
+    operator torch::Tensor() {
+      if (value.is_nil()) {
+        return {};
+      }
+      return from_ruby<torch::Tensor>(value);
+    }
+};
+
+template<>
+inline
+OptionalTensor from_ruby<OptionalTensor>(Object x)
+{
+  return OptionalTensor(x);
+}
+
+class ScalarType {
+  Object value;
+  public:
+    ScalarType(Object o) {
+      value = o;
+    }
+    operator at::ScalarType() {
+      throw std::runtime_error("ScalarType arguments not implemented yet");
+    }
+};
+
+template<>
+inline
+ScalarType from_ruby<ScalarType>(Object x)
+{
+  return ScalarType(x);
+}
+
+class OptionalScalarType {
+  Object value;
+  public:
+    OptionalScalarType(Object o) {
+      value = o;
+    }
+    operator c10::optional<at::ScalarType>() {
+      if (value.is_nil()) {
+        return c10::nullopt;
+      }
+      return ScalarType(value);
+    }
+};
+
+template<>
+inline
+OptionalScalarType from_ruby<OptionalScalarType>(Object x)
+{
+  return OptionalScalarType(x);
+}
+
+typedef torch::Device Device;
+
+Object wrap(std::tuple<torch::Tensor, torch::Tensor> x);
+Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> x);
+Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> x);
+Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> x);
+Object wrap(std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, int64_t> x);
+Object wrap(std::tuple<torch::Tensor, torch::Tensor, double, int64_t> x);

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"
@@ -22,31 +27,145 @@ require "torch/optim/sgd"
 require "torch/optim/lr_scheduler/lr_scheduler"
 require "torch/optim/lr_scheduler/step_lr"
 
-# nn base classes
+# nn parameters
+require "torch/nn/parameter"
+require "torch/nn/utils"
+
+# nn containers
 require "torch/nn/module"
+require "torch/nn/sequential"
+
+# nn convolution layers
 require "torch/nn/convnd"
-require "torch/nn/dropoutnd"
+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"
+require "torch/nn/identity"
+require "torch/nn/linear"
 
-# nn
+# nn dropout layers
+require "torch/nn/dropoutnd"
 require "torch/nn/alpha_dropout"
-require "torch/nn/conv2d"
 require "torch/nn/dropout"
 require "torch/nn/dropout2d"
 require "torch/nn/dropout3d"
-require "torch/nn/embedding"
 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"
+require "torch/nn/softmax"
+require "torch/nn/softmax2d"
+require "torch/nn/softmin"
+
+# nn sparse layers
+require "torch/nn/embedding"
+require "torch/nn/embedding_bag"
+
+# nn distance functions
+require "torch/nn/cosine_similarity"
+require "torch/nn/pairwise_distance"
+
+# nn loss functions
+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/cross_entropy_loss"
+require "torch/nn/ctc_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/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/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"
+
+# nn other
 require "torch/nn/functional"
 require "torch/nn/init"
-require "torch/nn/linear"
-require "torch/nn/mse_loss"
-require "torch/nn/parameter"
-require "torch/nn/relu"
-require "torch/nn/sequential"
 
 # utils
 require "torch/utils/data/data_loader"
 require "torch/utils/data/tensor_dataset"
 
+# random
+require "torch/random"
+
 module Torch
   class Error < StandardError; end
   class NotImplementedYet < StandardError
@@ -57,7 +176,6 @@ module Torch
 
   # keys: https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.dtype
   # values: https://github.com/pytorch/pytorch/blob/master/c10/core/ScalarType.h
-  # complex and quantized types not supported by PyTorch yet
   DTYPE_TO_ENUM = {
     uint8: 0,
     int8: 1,
@@ -73,17 +191,52 @@ module Torch
     float32: 6,
     double: 7,
     float64: 7,
-    # complex_half: 8,
-    # complex_float: 9,
-    # complex_double: 10,
+    complex_half: 8,
+    complex_float: 9,
+    complex_double: 10,
     bool: 11,
-    # qint8: 12,
-    # quint8: 13,
-    # qint32: 14,
-    # bfloat16: 15
+    qint8: 12,
+    quint8: 13,
+    qint32: 14,
+    bfloat16: 15
   }
   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|
@@ -120,6 +273,8 @@ module Torch
     # use method for cases when Numo not available
     # or available after Torch loaded
     def _dtype_to_numo
+      raise Error, "Numo not found" unless defined?(Numo::NArray)
+
       {
         uint8: Numo::UInt8,
         int8: Numo::Int8,
@@ -131,6 +286,29 @@ 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
+
+    def save(obj, f)
+      raise NotImplementedYet unless obj.is_a?(Tensor)
+      File.binwrite(f, _save(obj))
+    end
+
+    def load(f)
+      raise NotImplementedYet
+    end
+
     # --- begin tensor creation: https://pytorch.org/cppdocs/notes/tensor_creation.html ---
 
     def arange(start, finish = nil, step = 1, **options)
@@ -200,8 +378,12 @@ module Torch
         data = [data].compact
       end
 
-      if options[:dtype].nil? && data.all? { |v| v.is_a?(Integer) }
-        options[:dtype] = :int64
+      if options[:dtype].nil?
+        if data.all? { |v| v.is_a?(Integer) }
+          options[:dtype] = :int64
+        elsif data.all? { |v| v == true || v == false }
+          options[:dtype] = :bool
+        end
       end
 
       _tensor(data, size, tensor_options(**options))
@@ -210,19 +392,19 @@ module Torch
     # --- begin like ---
 
     def ones_like(input, **options)
-      ones(input.size, like_options(input, options))
+      ones(input.size, **like_options(input, options))
     end
 
     def empty_like(input, **options)
-      empty(input.size, like_options(input, options))
+      empty(input.size, **like_options(input, options))
     end
 
     def full_like(input, fill_value, **options)
-      full(input.size, fill_value, like_options(input, options))
+      full(input.size, fill_value, **like_options(input, options))
     end
 
     def rand_like(input, **options)
-      rand(input.size, like_options(input, options))
+      rand(input.size, **like_options(input, options))
     end
 
     def randint_like(input, low, high = nil, **options)
@@ -231,163 +413,19 @@ module Torch
         high = low
         low = 0
       end
-      randint(low, high, input.size, like_options(input, options))
+      randint(low, high, input.size, **like_options(input, options))
     end
 
     def randn_like(input, **options)
-      randn(input.size, like_options(input, options))
+      randn(input.size, **like_options(input, options))
     end
 
     def zeros_like(input, **options)
-      zeros(input.size, like_options(input, options))
-    end
-
-    # --- 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
-        _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 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 min(input)
-      _min(input)
-    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
-
-    def exp(input)
-      _exp(input)
-    end
-
-    def log(input)
-      _log(input)
-    end
-
-    def sign(input)
-      _sign(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
-
-    def abs(input)
-      _abs(input)
-    end
-
-    def device(str)
-      Device.new(str)
+      zeros(input.size, **like_options(input, options))
     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