torch-rb 0.4.2 → 0.6.0

data/ext/torch/tensor.cpp

@@ -0,0 +1,307 @@
+#include <torch/torch.h>
+
+#include <rice/Constructor.hpp>
+#include <rice/Module.hpp>
+
+#include "tensor_functions.h"
+#include "ruby_arg_parser.h"
+#include "templates.h"
+#include "utils.h"
+
+using namespace Rice;
+using torch::indexing::TensorIndex;
+
+Class rb_cTensor;
+
+std::vector<TensorIndex> index_vector(Array a) {
+  Object obj;
+
+  std::vector<TensorIndex> indices;
+  indices.reserve(a.size());
+
+  for (long i = 0; i < a.size(); i++) {
+    obj = a[i];
+
+    if (obj.is_instance_of(rb_cInteger)) {
+      indices.push_back(from_ruby<int64_t>(obj));
+    } else if (obj.is_instance_of(rb_cRange)) {
+      torch::optional<int64_t> start_index = torch::nullopt;
+      torch::optional<int64_t> stop_index = torch::nullopt;
+
+      Object begin = obj.call("begin");
+      if (!begin.is_nil()) {
+        start_index = from_ruby<int64_t>(begin);
+      }
+
+      Object end = obj.call("end");
+      if (!end.is_nil()) {
+        stop_index = from_ruby<int64_t>(end);
+      }
+
+      Object exclude_end = obj.call("exclude_end?");
+      if (stop_index.has_value() && !exclude_end) {
+        if (stop_index.value() == -1) {
+          stop_index = torch::nullopt;
+        } else {
+          stop_index = stop_index.value() + 1;
+        }
+      }
+
+      indices.push_back(torch::indexing::Slice(start_index, stop_index));
+    } else if (obj.is_instance_of(rb_cTensor)) {
+      indices.push_back(from_ruby<Tensor>(obj));
+    } else if (obj.is_nil()) {
+      indices.push_back(torch::indexing::None);
+    } else if (obj == True || obj == False) {
+      indices.push_back(from_ruby<bool>(obj));
+    } else {
+      throw Exception(rb_eArgError, "Unsupported index type: %s", rb_obj_classname(obj));
+    }
+  }
+  return indices;
+}
+
+// hack (removes inputs argument)
+// https://github.com/pytorch/pytorch/commit/2e5bfa9824f549be69a28e4705a72b4cf8a4c519
+// TODO add support for inputs argument
+// _backward
+static VALUE tensor__backward(int argc, VALUE* argv, VALUE self_)
+{
+  HANDLE_TH_ERRORS
+  Tensor& self = from_ruby<Tensor&>(self_);
+  static RubyArgParser parser({
+    "_backward(Tensor? gradient=None, bool? retain_graph=None, bool create_graph=False)"
+  });
+  ParsedArgs<4> parsed_args;
+  auto _r = parser.parse(self_, argc, argv, parsed_args);
+  // _backward(Tensor self, Tensor[] inputs, Tensor? gradient=None, bool? retain_graph=None, bool create_graph=False) -> ()
+  auto dispatch__backward = [](const Tensor & self, TensorList inputs, const OptionalTensor & gradient, c10::optional<bool> retain_graph, bool create_graph) -> void {
+    // in future, release GVL
+    self._backward(inputs, gradient, retain_graph, create_graph);
+  };
+  dispatch__backward(self, {}, _r.optionalTensor(0), _r.toBoolOptional(1), _r.toBool(2));
+  RETURN_NIL
+  END_HANDLE_TH_ERRORS
+}
+
+void init_tensor(Rice::Module& m) {
+  rb_cTensor = Rice::define_class_under<torch::Tensor>(m, "Tensor");
+  rb_cTensor.add_handler<torch::Error>(handle_error);
+  add_tensor_functions(rb_cTensor);
+  THPVariableClass = rb_cTensor.value();
+
+  rb_define_method(rb_cTensor, "backward", (VALUE (*)(...)) tensor__backward, -1);
+
+  rb_cTensor
+    .define_method("cuda?", &torch::Tensor::is_cuda)
+    .define_method("sparse?", &torch::Tensor::is_sparse)
+    .define_method("quantized?", &torch::Tensor::is_quantized)
+    .define_method("dim", &torch::Tensor::dim)
+    .define_method("numel", &torch::Tensor::numel)
+    .define_method("element_size", &torch::Tensor::element_size)
+    .define_method("requires_grad", &torch::Tensor::requires_grad)
+    .define_method(
+      "_size",
+      *[](Tensor& self, int64_t dim) {
+        return self.size(dim);
+      })
+    .define_method(
+      "_stride",
+      *[](Tensor& self, int64_t dim) {
+        return self.stride(dim);
+      })
+    // in C++ for performance
+    .define_method(
+      "shape",
+      *[](Tensor& self) {
+        Array a;
+        for (auto &size : self.sizes()) {
+          a.push(size);
+        }
+        return a;
+      })
+    .define_method(
+      "_strides",
+      *[](Tensor& self) {
+        Array a;
+        for (auto &stride : self.strides()) {
+          a.push(stride);
+        }
+        return a;
+      })
+    .define_method(
+      "_index",
+      *[](Tensor& self, Array indices) {
+        auto vec = index_vector(indices);
+        return self.index(vec);
+      })
+    .define_method(
+      "_index_put_custom",
+      *[](Tensor& self, Array indices, torch::Tensor& value) {
+        auto vec = index_vector(indices);
+        return self.index_put_(vec, value);
+      })
+    .define_method(
+      "contiguous?",
+      *[](Tensor& self) {
+        return self.is_contiguous();
+      })
+    .define_method(
+      "_requires_grad!",
+      *[](Tensor& self, bool requires_grad) {
+        return self.set_requires_grad(requires_grad);
+      })
+    .define_method(
+      "grad",
+      *[](Tensor& self) {
+        auto grad = self.grad();
+        return grad.defined() ? to_ruby<torch::Tensor>(grad) : Nil;
+      })
+    .define_method(
+      "grad=",
+      *[](Tensor& self, torch::Tensor& grad) {
+        self.mutable_grad() = grad;
+      })
+    .define_method(
+      "_dtype",
+      *[](Tensor& self) {
+        return (int) at::typeMetaToScalarType(self.dtype());
+      })
+    .define_method(
+      "_type",
+      *[](Tensor& self, int dtype) {
+        return self.toType((torch::ScalarType) dtype);
+      })
+    .define_method(
+      "_layout",
+      *[](Tensor& self) {
+        std::stringstream s;
+        s << self.layout();
+        return s.str();
+      })
+    .define_method(
+      "device",
+      *[](Tensor& self) {
+        std::stringstream s;
+        s << self.device();
+        return s.str();
+      })
+    .define_method(
+      "_data_str",
+      *[](Tensor& self) {
+        Tensor tensor = self;
+
+        // move to CPU to get data
+        if (tensor.device().type() != torch::kCPU) {
+          torch::Device device("cpu");
+          tensor = tensor.to(device);
+        }
+
+        if (!tensor.is_contiguous()) {
+          tensor = tensor.contiguous();
+        }
+
+        auto data_ptr = (const char *) tensor.data_ptr();
+        return std::string(data_ptr, tensor.numel() * tensor.element_size());
+      })
+    // for TorchVision
+    .define_method(
+      "_data_ptr",
+      *[](Tensor& self) {
+        return reinterpret_cast<uintptr_t>(self.data_ptr());
+      })
+    // TODO figure out a better way to do this
+    .define_method(
+      "_flat_data",
+      *[](Tensor& self) {
+        Tensor tensor = self;
+
+        // move to CPU to get data
+        if (tensor.device().type() != torch::kCPU) {
+          torch::Device device("cpu");
+          tensor = tensor.to(device);
+        }
+
+        Array a;
+        auto dtype = tensor.dtype();
+
+        Tensor view = tensor.reshape({tensor.numel()});
+
+        // TODO DRY if someone knows C++
+        if (dtype == torch::kByte) {
+          for (int i = 0; i < tensor.numel(); i++) {
+            a.push(view[i].item().to<uint8_t>());
+          }
+        } else if (dtype == torch::kChar) {
+          for (int i = 0; i < tensor.numel(); i++) {
+            a.push(to_ruby<int>(view[i].item().to<int8_t>()));
+          }
+        } else if (dtype == torch::kShort) {
+          for (int i = 0; i < tensor.numel(); i++) {
+            a.push(view[i].item().to<int16_t>());
+          }
+        } else if (dtype == torch::kInt) {
+          for (int i = 0; i < tensor.numel(); i++) {
+            a.push(view[i].item().to<int32_t>());
+          }
+        } else if (dtype == torch::kLong) {
+          for (int i = 0; i < tensor.numel(); i++) {
+            a.push(view[i].item().to<int64_t>());
+          }
+        } else if (dtype == torch::kFloat) {
+          for (int i = 0; i < tensor.numel(); i++) {
+            a.push(view[i].item().to<float>());
+          }
+        } else if (dtype == torch::kDouble) {
+          for (int i = 0; i < tensor.numel(); i++) {
+            a.push(view[i].item().to<double>());
+          }
+        } else if (dtype == torch::kBool) {
+          for (int i = 0; i < tensor.numel(); i++) {
+            a.push(view[i].item().to<bool>() ? True : False);
+          }
+        } else {
+          throw std::runtime_error("Unsupported type");
+        }
+        return a;
+      })
+    .define_method(
+      "_to",
+      *[](Tensor& self, torch::Device device, int dtype, bool non_blocking, bool copy) {
+        return self.to(device, (torch::ScalarType) dtype, non_blocking, copy);
+      });
+
+  Rice::define_class_under<torch::TensorOptions>(m, "TensorOptions")
+    .add_handler<torch::Error>(handle_error)
+    .define_constructor(Rice::Constructor<torch::TensorOptions>())
+    .define_method(
+      "dtype",
+      *[](torch::TensorOptions& self, int dtype) {
+        return self.dtype((torch::ScalarType) dtype);
+      })
+    .define_method(
+      "layout",
+      *[](torch::TensorOptions& self, const std::string& layout) {
+        torch::Layout l;
+        if (layout == "strided") {
+          l = torch::kStrided;
+        } else if (layout == "sparse") {
+          l = torch::kSparse;
+          throw std::runtime_error("Sparse layout not supported yet");
+        } else {
+          throw std::runtime_error("Unsupported layout: " + layout);
+        }
+        return self.layout(l);
+      })
+    .define_method(
+      "device",
+      *[](torch::TensorOptions& self, const std::string& device) {
+        torch::Device d(device);
+        return self.device(d);
+      })
+    .define_method(
+      "requires_grad",
+      *[](torch::TensorOptions& self, bool requires_grad) {
+        return self.requires_grad(requires_grad);
+      });
+}

data/ext/torch/tensor_functions.h

@@ -3,4 +3,4 @@
 
 #pragma once
 
-void add_tensor_functions(Module m);
+void add_tensor_functions(Rice::Module& m);

data/ext/torch/torch.cpp

@@ -0,0 +1,86 @@
+#include <torch/torch.h>
+
+#include <rice/Module.hpp>
+
+#include "torch_functions.h"
+#include "templates.h"
+#include "utils.h"
+
+void init_torch(Rice::Module& m) {
+  m.add_handler<torch::Error>(handle_error);
+  add_torch_functions(m);
+  m.define_singleton_method(
+      "grad_enabled?",
+      *[]() {
+        return torch::GradMode::is_enabled();
+      })
+    .define_singleton_method(
+      "_set_grad_enabled",
+      *[](bool enabled) {
+        torch::GradMode::set_enabled(enabled);
+      })
+    .define_singleton_method(
+      "manual_seed",
+      *[](uint64_t seed) {
+        return torch::manual_seed(seed);
+      })
+    // config
+    .define_singleton_method(
+      "show_config",
+      *[] {
+        return torch::show_config();
+      })
+    .define_singleton_method(
+      "parallel_info",
+      *[] {
+        return torch::get_parallel_info();
+      })
+    // begin operations
+    .define_singleton_method(
+      "_save",
+      *[](const torch::IValue &value) {
+        auto v = torch::pickle_save(value);
+        std::string str(v.begin(), v.end());
+        return str;
+      })
+    .define_singleton_method(
+      "_load",
+      *[](const std::string &s) {
+        std::vector<char> v;
+        std::copy(s.begin(), s.end(), std::back_inserter(v));
+        // https://github.com/pytorch/pytorch/issues/20356#issuecomment-567663701
+        return torch::pickle_load(v);
+      })
+    .define_singleton_method(
+      "_from_blob",
+      *[](Rice::String s, std::vector<int64_t> size, const torch::TensorOptions &options) {
+        void *data = const_cast<char *>(s.c_str());
+        return torch::from_blob(data, size, options);
+      })
+    .define_singleton_method(
+      "_tensor",
+      *[](Rice::Array a, std::vector<int64_t> size, const torch::TensorOptions &options) {
+        auto dtype = options.dtype();
+        torch::Tensor t;
+        if (dtype == torch::kBool) {
+          std::vector<uint8_t> vec;
+          for (long i = 0; i < a.size(); i++) {
+            vec.push_back(from_ruby<bool>(a[i]));
+          }
+          t = torch::tensor(vec, options);
+        } else {
+          std::vector<float> vec;
+          for (long i = 0; i < a.size(); i++) {
+            vec.push_back(from_ruby<float>(a[i]));
+          }
+          // hack for requires_grad error
+          if (options.requires_grad()) {
+            t = torch::tensor(vec, options.requires_grad(c10::nullopt));
+            t.set_requires_grad(true);
+          } else {
+            t = torch::tensor(vec, options);
+          }
+        }
+        return t.reshape(size);
+      });
+}

data/ext/torch/torch_functions.h

@@ -3,4 +3,4 @@
 
 #pragma once
 
-void add_torch_functions(Module m);
+void add_torch_functions(Rice::Module& m);

data/ext/torch/utils.h

@@ -1,13 +1,20 @@
 #pragma once
 
+#include <rice/Exception.hpp>
 #include <rice/Symbol.hpp>
 
+// TODO find better place
+inline void handle_error(torch::Error const & ex)
+{
+  throw Rice::Exception(rb_eRuntimeError, ex.what_without_backtrace());
+}
+
 // keep THP prefix for now to make it easier to compare code
 
 extern VALUE THPVariableClass;
 
 inline VALUE THPUtils_internSymbol(const std::string& str) {
-  return Symbol(str);
+  return Rice::Symbol(str);
 }
 
 inline std::string THPUtils_unpackSymbol(VALUE obj) {

data/ext/torch/wrap_outputs.h

@@ -90,3 +90,10 @@ inline Object wrap(torch::TensorList x) {
   }
   return Object(a);
 }
+
+inline Object wrap(std::tuple<double, double> x) {
+  Array a;
+  a.push(to_ruby<double>(std::get<0>(x)));
+  a.push(to_ruby<double>(std::get<1>(x)));
+  return Object(a);
+}

data/lib/torch.rb

@@ -261,6 +261,8 @@ module Torch
         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)
+      elsif args.size == 0
+        Torch.empty(0, dtype: dtype, device: device)
       else
         Torch.empty(*args, dtype: dtype, device: device)
       end
@@ -335,25 +337,24 @@ module Torch
       }
     end
 
-    def no_grad
-      previous_value = grad_enabled?
-      begin
-        _set_grad_enabled(false)
-        yield
-      ensure
-        _set_grad_enabled(previous_value)
-      end
+    def no_grad(&block)
+      grad_enabled(false, &block)
+    end
+
+    def enable_grad(&block)
+      grad_enabled(true, &block)
     end
 
-    def enable_grad
+    def grad_enabled(value)
       previous_value = grad_enabled?
       begin
-        _set_grad_enabled(true)
+        _set_grad_enabled(value)
         yield
       ensure
         _set_grad_enabled(previous_value)
       end
     end
+    alias_method :set_grad_enabled, :grad_enabled
 
     def device(str)
       Device.new(str)
@@ -434,7 +435,8 @@ 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)
+    # center option
+    def stft(input, n_fft, hop_length: nil, win_length: nil, window: nil, center: true, pad_mode: "reflect", normalized: false, onesided: true, return_complex: nil)
       if center
         signal_dim = input.dim
         extended_shape = [1] * (3 - signal_dim) + input.size
@@ -442,12 +444,7 @@ module Torch
         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)
+      _stft(input, n_fft, hop_length, win_length, window, normalized, onesided, return_complex)
     end
 
     private