torch-rb 0.5.0 → 0.7.0

data/ext/torch/templates.h

@@ -4,8 +4,7 @@
 #undef isfinite
 #endif
 
-#include <rice/Array.hpp>
-#include <rice/Object.hpp>
+#include <rice/rice.hpp>
 
 using namespace Rice;
 
@@ -23,6 +22,9 @@ using torch::ArrayRef;
 using torch::TensorList;
 using torch::Storage;
 
+using torch::nn::init::FanModeType;
+using torch::nn::init::NonlinearityType;
+
 #define HANDLE_TH_ERRORS                                             \
   try {
 
@@ -38,37 +40,42 @@ using torch::Storage;
 #define RETURN_NIL                                                   \
   return Qnil;
 
-template<>
-inline
-std::vector<int64_t> from_ruby<std::vector<int64_t>>(Object x)
-{
-  Array a = Array(x);
-  std::vector<int64_t> vec(a.size());
-  for (size_t i = 0; i < a.size(); i++) {
-    vec[i] = from_ruby<int64_t>(a[i]);
-  }
-  return vec;
-}
+class OptionalTensor {
+  torch::Tensor value;
+  public:
+    OptionalTensor(Object o) {
+      if (o.is_nil()) {
+        value = {};
+      } else {
+        value = Rice::detail::From_Ruby<torch::Tensor>().convert(o.value());
+      }
+    }
+    OptionalTensor(torch::Tensor o) {
+      value = o;
+    }
+    operator torch::Tensor() const {
+      return value;
+    }
+};
 
-template<>
-inline
-std::vector<Tensor> from_ruby<std::vector<Tensor>>(Object x)
+namespace Rice::detail
 {
-  Array a = Array(x);
-  std::vector<Tensor> vec(a.size());
-  for (size_t i = 0; i < a.size(); i++) {
-    vec[i] = from_ruby<Tensor>(a[i]);
-  }
-  return vec;
-}
+  template<>
+  struct Type<FanModeType>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
 
-class FanModeType {
-  std::string s;
+  template<>
+  class From_Ruby<FanModeType>
+  {
   public:
-    FanModeType(Object o) {
-      s = String(o).str();
-    }
-    operator torch::nn::init::FanModeType() {
+    FanModeType convert(VALUE x)
+    {
+      auto s = String(x).str();
       if (s == "fan_in") {
         return torch::kFanIn;
       } else if (s == "fan_out") {
@@ -77,22 +84,24 @@ class FanModeType {
         throw std::runtime_error("Unsupported nonlinearity type: " + s);
       }
     }
-};
-
-template<>
-inline
-FanModeType from_ruby<FanModeType>(Object x)
-{
-  return FanModeType(x);
-}
+  };
+
+  template<>
+  struct Type<NonlinearityType>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
 
-class NonlinearityType {
-  std::string s;
+  template<>
+  class From_Ruby<NonlinearityType>
+  {
   public:
-    NonlinearityType(Object o) {
-      s = String(o).str();
-    }
-    operator torch::nn::init::NonlinearityType() {
+    NonlinearityType convert(VALUE x)
+    {
+      auto s = String(x).str();
       if (s == "linear") {
         return torch::kLinear;
       } else if (s == "conv1d") {
@@ -119,102 +128,70 @@ class NonlinearityType {
         throw std::runtime_error("Unsupported nonlinearity type: " + s);
       }
     }
-};
+  };
+
+  template<>
+  struct Type<OptionalTensor>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
 
-template<>
-inline
-NonlinearityType from_ruby<NonlinearityType>(Object x)
-{
-  return NonlinearityType(x);
-}
+  template<>
+  class From_Ruby<OptionalTensor>
+  {
+  public:
+    OptionalTensor convert(VALUE x)
+    {
+      return OptionalTensor(x);
+    }
+  };
+
+  template<>
+  struct Type<Scalar>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
 
-class OptionalTensor {
-  torch::Tensor value;
+  template<>
+  class From_Ruby<Scalar>
+  {
   public:
-    OptionalTensor(Object o) {
-      if (o.is_nil()) {
-        value = {};
+    Scalar convert(VALUE x)
+    {
+      if (FIXNUM_P(x)) {
+        return torch::Scalar(From_Ruby<int64_t>().convert(x));
       } else {
-        value = from_ruby<torch::Tensor>(o);
+        return torch::Scalar(From_Ruby<double>().convert(x));
       }
     }
-    OptionalTensor(torch::Tensor o) {
-      value = o;
-    }
-    operator torch::Tensor() const {
-      return value;
+  };
+
+  template<typename T>
+  struct Type<torch::optional<T>>
+  {
+    static bool verify()
+    {
+      return true;
     }
-};
-
-template<>
-inline
-Scalar from_ruby<Scalar>(Object x)
-{
-  if (x.rb_type() == T_FIXNUM) {
-    return torch::Scalar(from_ruby<int64_t>(x));
-  } else {
-    return torch::Scalar(from_ruby<double>(x));
-  }
-}
-
-template<>
-inline
-OptionalTensor from_ruby<OptionalTensor>(Object x)
-{
-  return OptionalTensor(x);
-}
-
-template<>
-inline
-torch::optional<torch::ScalarType> from_ruby<torch::optional<torch::ScalarType>>(Object x)
-{
-  if (x.is_nil()) {
-    return torch::nullopt;
-  } else {
-    return torch::optional<torch::ScalarType>{from_ruby<torch::ScalarType>(x)};
-  }
-}
-
-template<>
-inline
-torch::optional<int64_t> from_ruby<torch::optional<int64_t>>(Object x)
-{
-  if (x.is_nil()) {
-    return torch::nullopt;
-  } else {
-    return torch::optional<int64_t>{from_ruby<int64_t>(x)};
-  }
-}
+  };
 
-template<>
-inline
-torch::optional<double> from_ruby<torch::optional<double>>(Object x)
-{
-  if (x.is_nil()) {
-    return torch::nullopt;
-  } else {
-    return torch::optional<double>{from_ruby<double>(x)};
-  }
-}
-
-template<>
-inline
-torch::optional<bool> from_ruby<torch::optional<bool>>(Object x)
-{
-  if (x.is_nil()) {
-    return torch::nullopt;
-  } else {
-    return torch::optional<bool>{from_ruby<bool>(x)};
-  }
-}
-
-template<>
-inline
-torch::optional<Scalar> from_ruby<torch::optional<Scalar>>(Object x)
-{
-  if (x.is_nil()) {
-    return torch::nullopt;
-  } else {
-    return torch::optional<Scalar>{from_ruby<Scalar>(x)};
-  }
+  template<typename T>
+  class From_Ruby<torch::optional<T>>
+  {
+  public:
+    torch::optional<T> convert(VALUE x)
+    {
+      if (NIL_P(x)) {
+        return torch::nullopt;
+      } else {
+        return torch::optional<T>{From_Ruby<T>().convert(x)};
+      }
+    }
+  };
 }

data/ext/torch/tensor.cpp

@@ -0,0 +1,320 @@
+#include <torch/torch.h>
+
+#include <rice/rice.hpp>
+
+#include "tensor_functions.h"
+#include "ruby_arg_parser.h"
+#include "templates.h"
+#include "utils.h"
+
+using namespace Rice;
+using torch::indexing::TensorIndex;
+
+namespace Rice::detail
+{
+  template<typename T>
+  struct Type<c10::complex<T>>
+  {
+    static bool verify()
+    {
+      return true;
+    }
+  };
+
+  template<typename T>
+  class To_Ruby<c10::complex<T>>
+  {
+  public:
+    VALUE convert(c10::complex<T> const& x)
+    {
+      return rb_dbl_complex_new(x.real(), x.imag());
+    }
+  };
+}
+
+template<typename T>
+Array flat_data(Tensor& tensor) {
+  Tensor view = tensor.reshape({tensor.numel()});
+
+  Array a;
+  for (int i = 0; i < tensor.numel(); i++) {
+    a.push(view[i].item().to<T>());
+  }
+  return a;
+}
+
+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(Rice::detail::From_Ruby<int64_t>().convert(obj.value()));
+    } 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 = Rice::detail::From_Ruby<int64_t>().convert(begin.value());
+      }
+
+      Object end = obj.call("end");
+      if (!end.is_nil()) {
+        stop_index = Rice::detail::From_Ruby<int64_t>().convert(end.value());
+      }
+
+      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(Rice::detail::From_Ruby<Tensor>().convert(obj.value()));
+    } else if (obj.is_nil()) {
+      indices.push_back(torch::indexing::None);
+    } else if (obj == True || obj == False) {
+      indices.push_back(Rice::detail::From_Ruby<bool>().convert(obj.value()));
+    } 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 = Rice::detail::From_Ruby<Tensor&>().convert(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, Rice::Class& c, Rice::Class& rb_cTensorOptions) {
+  rb_cTensor = c;
+  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() ? Object(Rice::detail::To_Ruby<torch::Tensor>().convert(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) {
+        auto 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) {
+        auto tensor = self;
+
+        // move to CPU to get data
+        if (tensor.device().type() != torch::kCPU) {
+          torch::Device device("cpu");
+          tensor = tensor.to(device);
+        }
+
+        auto dtype = tensor.dtype();
+        if (dtype == torch::kByte) {
+          return flat_data<uint8_t>(tensor);
+        } else if (dtype == torch::kChar) {
+          return flat_data<int8_t>(tensor);
+        } else if (dtype == torch::kShort) {
+          return flat_data<int16_t>(tensor);
+        } else if (dtype == torch::kInt) {
+          return flat_data<int32_t>(tensor);
+        } else if (dtype == torch::kLong) {
+          return flat_data<int64_t>(tensor);
+        } else if (dtype == torch::kFloat) {
+          return flat_data<float>(tensor);
+        } else if (dtype == torch::kDouble) {
+          return flat_data<double>(tensor);
+        } else if (dtype == torch::kBool) {
+          return flat_data<bool>(tensor);
+        } else if (dtype == torch::kComplexFloat) {
+          return flat_data<c10::complex<float>>(tensor);
+        } else if (dtype == torch::kComplexDouble) {
+          return flat_data<c10::complex<double>>(tensor);
+        } else {
+          throw std::runtime_error("Unsupported type");
+        }
+      })
+    .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);
+      });
+
+  rb_cTensorOptions
+    .add_handler<torch::Error>(handle_error)
+    .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);
+      });
+}