torch-rb 0.3.4 → 0.4.1

data/ext/torch/ruby_arg_parser.h

@@ -0,0 +1,373 @@
+// adapted from PyTorch - python_arg_parser.h
+
+#pragma once
+
+#include <torch/torch.h>
+#include <rice/Exception.hpp>
+
+#include "templates.h"
+#include "utils.h"
+
+enum class ParameterType {
+  TENSOR, SCALAR, INT64, DOUBLE, COMPLEX, TENSOR_LIST, INT_LIST, GENERATOR,
+  BOOL, STORAGE, PYOBJECT, SCALARTYPE, LAYOUT, MEMORY_FORMAT, DEVICE, STRING,
+  DIMNAME, DIMNAME_LIST, QSCHEME, FLOAT_LIST
+};
+
+struct FunctionParameter {
+  FunctionParameter(const std::string& fmt, bool keyword_only);
+
+  bool check(VALUE obj, int argnum);
+
+  void set_default_str(const std::string& str);
+  std::string type_name() const;
+
+  ParameterType type_;
+  bool optional;
+  bool allow_none;
+  bool keyword_only;
+  bool allow_numbers_as_tensors = false;
+  int size;
+  std::string name;
+  VALUE ruby_name;
+  at::SmallVector<VALUE, 5> numpy_python_names;
+  at::Scalar default_scalar;
+  std::vector<int64_t> default_intlist;
+  union {
+    bool default_bool;
+    int64_t default_int;
+    double default_double;
+    double default_complex[2]; // see Scalar
+    at::ScalarType default_scalartype;
+    at::Layout default_layout;
+  };
+};
+
+struct FunctionSignature {
+  explicit FunctionSignature(const std::string& fmt, int index);
+
+  bool parse(VALUE self, VALUE args, VALUE kwargs, std::vector<VALUE>& dst, bool raise_exception);
+
+  std::string toString() const;
+
+  std::string name;
+  std::vector<FunctionParameter> params;
+  // std::vector<py::handle> overloaded_args;
+  ssize_t min_args;
+  ssize_t max_args;
+  ssize_t max_pos_args;
+  int index;
+  bool hidden;
+  bool deprecated;
+  bool disable_torch_function;
+};
+
+struct RubyArgs {
+  RubyArgs(const FunctionSignature& signature, std::vector<VALUE> &args)
+    : signature(signature)
+    , args(args)
+    , idx(signature.index) {}
+
+  const FunctionSignature& signature;
+  std::vector<VALUE> args;
+  int idx;
+
+  inline at::Tensor tensor(int i);
+  inline OptionalTensor optionalTensor(int i);
+  inline at::Scalar scalar(int i);
+  // inline at::Scalar scalarWithDefault(int i, at::Scalar default_scalar);
+  inline std::vector<at::Tensor> tensorlist(int i);
+  template<int N>
+  inline std::array<at::Tensor, N> tensorlist_n(int i);
+  inline std::vector<int64_t> intlist(int i);
+  // inline c10::OptionalArray<int64_t> intlistOptional(int i);
+  // inline std::vector<int64_t> intlistWithDefault(int i, std::vector<int64_t> default_intlist);
+  inline c10::optional<at::Generator> generator(int i);
+  inline at::Storage storage(int i);
+  inline at::ScalarType scalartype(int i);
+  // inline at::ScalarType scalartypeWithDefault(int i, at::ScalarType default_scalartype);
+  inline c10::optional<at::ScalarType> scalartypeOptional(int i);
+  inline c10::optional<at::Scalar> scalarOptional(int i);
+  inline c10::optional<int64_t> toInt64Optional(int i);
+  inline c10::optional<bool> toBoolOptional(int i);
+  inline c10::optional<double> toDoubleOptional(int i);
+  // inline c10::OptionalArray<double> doublelistOptional(int i);
+  // inline at::Layout layout(int i);
+  // inline at::Layout layoutWithDefault(int i, at::Layout default_layout);
+  inline c10::optional<at::Layout> layoutOptional(int i);
+  inline at::Device device(int i);
+  // inline at::Device deviceWithDefault(int i, const at::Device& default_device);
+  // inline c10::optional<at::Device> deviceOptional(int i);
+  // inline at::Dimname dimname(int i);
+  // inline std::vector<at::Dimname> dimnamelist(int i);
+  // inline c10::optional<std::vector<at::Dimname>> toDimnameListOptional(int i);
+  inline at::MemoryFormat memoryformat(int i);
+  inline c10::optional<at::MemoryFormat> memoryformatOptional(int i);
+  // inline at::QScheme toQScheme(int i);
+  inline std::string string(int i);
+  // inline c10::optional<std::string> stringOptional(int i);
+  // inline PyObject* pyobject(int i);
+  inline int64_t toInt64(int i);
+  // inline int64_t toInt64WithDefault(int i, int64_t default_int);
+  inline double toDouble(int i);
+  // inline double toDoubleWithDefault(int i, double default_double);
+  // inline c10::complex<double> toComplex(int i);
+  // inline c10::complex<double> toComplexWithDefault(int i, c10::complex<double> default_complex);
+  inline bool toBool(int i);
+  // inline bool toBoolWithDefault(int i, bool default_bool);
+  inline bool isNone(int i);
+};
+
+inline at::Tensor RubyArgs::tensor(int i) {
+  return from_ruby<torch::Tensor>(args[i]);
+}
+
+inline OptionalTensor RubyArgs::optionalTensor(int i) {
+  if (NIL_P(args[i])) return OptionalTensor(Nil);
+  return tensor(i);
+}
+
+inline at::Scalar RubyArgs::scalar(int i) {
+  if (NIL_P(args[i])) return signature.params[i].default_scalar;
+  return from_ruby<torch::Scalar>(args[i]);
+}
+
+inline std::vector<at::Tensor> RubyArgs::tensorlist(int i) {
+  if (NIL_P(args[i])) return std::vector<at::Tensor>();
+  return from_ruby<std::vector<Tensor>>(args[i]);
+}
+
+template<int N>
+inline std::array<at::Tensor, N> RubyArgs::tensorlist_n(int i) {
+  auto res = std::array<at::Tensor, N>();
+  if (NIL_P(args[i])) return res;
+  VALUE arg = args[i];
+  Check_Type(arg, T_ARRAY);
+  auto size = RARRAY_LEN(arg);
+  if (size != N) {
+    rb_raise(rb_eArgError, "expected array of %d elements but got %d", N, (int)size);
+  }
+  for (int idx = 0; idx < size; idx++) {
+    VALUE obj = rb_ary_entry(arg, idx);
+    res[idx] = from_ruby<Tensor>(obj);
+  }
+  return res;
+}
+
+inline std::vector<int64_t> RubyArgs::intlist(int i) {
+  if (NIL_P(args[i])) return signature.params[i].default_intlist;
+
+  VALUE arg = args[i];
+  auto size = signature.params[i].size;
+  if (size > 0 && FIXNUM_P(arg)) {
+    return std::vector<int64_t>(size, FIX2INT(arg));
+  }
+
+  size = RARRAY_LEN(arg);
+  std::vector<int64_t> res(size);
+  for (idx = 0; idx < size; idx++) {
+    VALUE obj = rb_ary_entry(arg, idx);
+    if (FIXNUM_P(obj)) {
+      res[idx] = from_ruby<int64_t>(obj);
+    } else {
+      rb_raise(rb_eArgError, "%s(): argument '%s' must be %s, but found element of type %s at pos %d",
+          signature.name.c_str(), signature.params[i].name.c_str(),
+          signature.params[i].type_name().c_str(), rb_obj_classname(obj), idx + 1);
+    }
+  }
+  return res;
+}
+
+inline c10::optional<at::Generator> RubyArgs::generator(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
+  throw std::runtime_error("generator not supported yet");
+}
+
+inline at::Storage RubyArgs::storage(int i) {
+  if (NIL_P(args[i])) return at::Storage();
+  throw std::runtime_error("storage not supported yet");
+}
+
+inline ScalarType RubyArgs::scalartype(int i) {
+  if (NIL_P(args[i])) {
+    auto scalartype = signature.params[i].default_scalartype;
+    return (scalartype == at::ScalarType::Undefined) ? at::typeMetaToScalarType(at::get_default_dtype()) : scalartype;
+  }
+
+  static std::unordered_map<VALUE, ScalarType> dtype_map = {
+    {ID2SYM(rb_intern("uint8")), ScalarType::Byte},
+    {ID2SYM(rb_intern("int8")), ScalarType::Char},
+    {ID2SYM(rb_intern("short")), ScalarType::Short},
+    {ID2SYM(rb_intern("int16")), ScalarType::Short},
+    {ID2SYM(rb_intern("int")), ScalarType::Int},
+    {ID2SYM(rb_intern("int32")), ScalarType::Int},
+    {ID2SYM(rb_intern("long")), ScalarType::Long},
+    {ID2SYM(rb_intern("int64")), ScalarType::Long},
+    {ID2SYM(rb_intern("float")), ScalarType::Float},
+    {ID2SYM(rb_intern("float32")), ScalarType::Float},
+    {ID2SYM(rb_intern("double")), ScalarType::Double},
+    {ID2SYM(rb_intern("float64")), ScalarType::Double},
+    {ID2SYM(rb_intern("complex_half")), ScalarType::ComplexHalf},
+    {ID2SYM(rb_intern("complex_float")), ScalarType::ComplexFloat},
+    {ID2SYM(rb_intern("complex_double")), ScalarType::ComplexDouble},
+    {ID2SYM(rb_intern("bool")), ScalarType::Bool},
+    {ID2SYM(rb_intern("qint8")), ScalarType::QInt8},
+    {ID2SYM(rb_intern("quint8")), ScalarType::QUInt8},
+    {ID2SYM(rb_intern("qint32")), ScalarType::QInt32},
+    {ID2SYM(rb_intern("bfloat16")), ScalarType::BFloat16},
+  };
+
+  auto it = dtype_map.find(args[i]);
+  if (it == dtype_map.end()) {
+    rb_raise(rb_eArgError, "invalid dtype: %s", THPUtils_unpackSymbol(args[i]).c_str());
+  }
+  return it->second;
+}
+
+inline c10::optional<ScalarType> RubyArgs::scalartypeOptional(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
+  return scalartype(i);
+}
+
+inline c10::optional<Scalar> RubyArgs::scalarOptional(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
+  return scalar(i);
+}
+
+inline c10::optional<int64_t> RubyArgs::toInt64Optional(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
+  return toInt64(i);
+}
+
+inline c10::optional<bool> RubyArgs::toBoolOptional(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
+  return toBool(i);
+}
+
+inline c10::optional<double> RubyArgs::toDoubleOptional(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
+  return toDouble(i);
+}
+
+inline c10::optional<at::Layout> RubyArgs::layoutOptional(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
+
+  static std::unordered_map<VALUE, Layout> layout_map = {
+    {ID2SYM(rb_intern("strided")), Layout::Strided},
+  };
+
+  auto it = layout_map.find(args[i]);
+  if (it == layout_map.end()) {
+    rb_raise(rb_eArgError, "invalid layout: %s", THPUtils_unpackSymbol(args[i]).c_str());
+  }
+  return it->second;
+}
+
+inline at::Device RubyArgs::device(int i) {
+  if (NIL_P(args[i])) {
+    return at::Device("cpu");
+  }
+  const std::string &device_str = THPUtils_unpackString(args[i]);
+  return at::Device(device_str);
+}
+
+inline at::MemoryFormat RubyArgs::memoryformat(int i) {
+  if (NIL_P(args[i])) return at::MemoryFormat::Contiguous;
+  throw std::runtime_error("memoryformat not supported yet");
+}
+
+inline c10::optional<at::MemoryFormat> RubyArgs::memoryformatOptional(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
+  return memoryformat(i);
+}
+
+inline std::string RubyArgs::string(int i) {
+  return from_ruby<std::string>(args[i]);
+}
+
+inline int64_t RubyArgs::toInt64(int i) {
+  if (NIL_P(args[i])) return signature.params[i].default_int;
+  return from_ruby<int64_t>(args[i]);
+}
+
+inline double RubyArgs::toDouble(int i) {
+  if (NIL_P(args[i])) return signature.params[i].default_double;
+  return from_ruby<double>(args[i]);
+}
+
+inline bool RubyArgs::toBool(int i) {
+  if (NIL_P(args[i])) return signature.params[i].default_bool;
+  return RTEST(args[i]);
+}
+
+inline bool RubyArgs::isNone(int i) {
+  return NIL_P(args[i]);
+}
+
+struct RubyArgParser {
+  std::vector<FunctionSignature> signatures_;
+  std::string function_name;
+  ssize_t max_args;
+
+  public:
+    RubyArgParser(std::vector<std::string> fmts) : max_args(0) {
+      int index = 0;
+      for (auto& fmt : fmts) {
+        signatures_.emplace_back(fmt, index);
+        ++index;
+      }
+      for (auto& signature : signatures_) {
+        if (signature.max_args > max_args) {
+          max_args = signature.max_args;
+        }
+      }
+      if (signatures_.size() > 0) {
+        function_name = signatures_[0].name;
+      }
+
+      // Check deprecated signatures last
+      std::stable_partition(signatures_.begin(), signatures_.end(),
+        [](const FunctionSignature & sig) {
+          return !sig.deprecated;
+        });
+    }
+
+    RubyArgs parse(VALUE self, int argc, VALUE* argv, std::vector<VALUE> &parsed_args) {
+      VALUE args, kwargs;
+      rb_scan_args(argc, argv, "*:", &args, &kwargs);
+
+      if (signatures_.size() == 1) {
+        auto& signature = signatures_[0];
+        signature.parse(self, args, kwargs, parsed_args, true);
+        return RubyArgs(signature, parsed_args);
+      }
+
+      for (auto& signature : signatures_) {
+        if (signature.parse(self, args, kwargs, parsed_args, false)) {
+          return RubyArgs(signature, parsed_args);
+        }
+      }
+
+      print_error(self, args, kwargs, parsed_args);
+
+      // TODO better message
+      rb_raise(rb_eArgError, "No matching signatures");
+    }
+
+    void print_error(VALUE self, VALUE args, VALUE kwargs, std::vector<VALUE>& parsed_args) {
+      ssize_t num_args = (NIL_P(args) ? 0 : RARRAY_LEN(args)) + (NIL_P(kwargs) ? 0 : RHASH_SIZE(kwargs));
+      std::vector<int> plausible_idxs;
+      ssize_t i = 0;
+      for (auto& signature : signatures_) {
+        if (num_args >= signature.min_args && num_args <= signature.max_args && !signature.hidden) {
+          plausible_idxs.push_back(i);
+        }
+        i++;
+      }
+
+      if (plausible_idxs.size() == 1) {
+        auto& signature = signatures_[plausible_idxs[0]];
+        signature.parse(self, args, kwargs, parsed_args, true);
+      }
+    }
+};

data/ext/torch/{templates.hpp → templates.h}

@@ -10,75 +10,55 @@
 using namespace Rice;
 
 using torch::Device;
+using torch::Scalar;
 using torch::ScalarType;
 using torch::Tensor;
+using torch::QScheme;
+using torch::Generator;
+using torch::TensorOptions;
+using torch::Layout;
+using torch::MemoryFormat;
+using torch::IntArrayRef;
+using torch::TensorList;
+using torch::Storage;
+
+#define HANDLE_TH_ERRORS                                             \
+  try {
+
+#define END_HANDLE_TH_ERRORS                                         \
+  } catch (const torch::Error& ex) {                                 \
+    rb_raise(rb_eRuntimeError, "%s", ex.what_without_backtrace());   \
+  } catch (const Rice::Exception& ex) {                              \
+    rb_raise(ex.class_of(), "%s", ex.what());                        \
+  } catch (const std::exception& ex) {                               \
+    rb_raise(rb_eRuntimeError, "%s", ex.what());                     \
+  }
 
-// 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;
-    }
-};
+#define RETURN_NIL                                                   \
+  return Qnil;
 
 template<>
 inline
-Scalar from_ruby<Scalar>(Object x)
+std::vector<int64_t> from_ruby<std::vector<int64_t>>(Object x)
 {
-  return Scalar(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 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)
+std::vector<Tensor> from_ruby<std::vector<Tensor>>(Object x)
 {
-  return TensorList(x);
+  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;
 }
 
 class FanModeType {
@@ -147,51 +127,35 @@ NonlinearityType from_ruby<NonlinearityType>(Object x)
   return NonlinearityType(x);
 }
 
-class MyReduction {
-  Object value;
+class OptionalTensor {
+  torch::Tensor 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 if (s == "none") {
-        return torch::Reduction::None;
+    OptionalTensor(Object o) {
+      if (o.is_nil()) {
+        value = {};
       } else {
-        throw std::runtime_error("Unsupported reduction: " + s);
+        value = from_ruby<torch::Tensor>(o);
       }
     }
+    OptionalTensor(torch::Tensor o) {
+      value = o;
+    }
+    operator torch::Tensor() const {
+      return value;
+    }
 };
 
 template<>
 inline
-MyReduction from_ruby<MyReduction>(Object x)
+Scalar from_ruby<Scalar>(Object x)
 {
-  return MyReduction(x);
+  if (x.rb_type() == T_FIXNUM) {
+    return torch::Scalar(from_ruby<int64_t>(x));
+  } else {
+    return torch::Scalar(from_ruby<double>(x));
+  }
 }
 
-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)
@@ -221,9 +185,35 @@ torch::optional<int64_t> from_ruby<torch::optional<int64_t>>(Object x)
   }
 }
 
-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);
+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)};
+  }
+}