torch-rb 0.11.2 → 0.12.1

data/ext/torch/ext.cpp

@@ -12,6 +12,7 @@ void init_torch(Rice::Module& m);
 void init_backends(Rice::Module& m);
 void init_cuda(Rice::Module& m);
 void init_device(Rice::Module& m);
+void init_generator(Rice::Module& m, Rice::Class& rb_cGenerator);
 void init_ivalue(Rice::Module& m, Rice::Class& rb_cIValue);
 void init_random(Rice::Module& m);
 
@@ -23,6 +24,7 @@ void Init_ext()
   // need to define certain classes up front to keep Rice happy
   auto rb_cIValue = Rice::define_class_under<torch::IValue>(m, "IValue")
     .define_constructor(Rice::Constructor<torch::IValue>());
+  auto rb_cGenerator = Rice::define_class_under<torch::Generator>(m, "Generator");
   auto rb_cTensor = Rice::define_class_under<torch::Tensor>(m, "Tensor");
   auto rb_cTensorOptions = Rice::define_class_under<torch::TensorOptions>(m, "TensorOptions")
     .define_constructor(Rice::Constructor<torch::TensorOptions>());
@@ -38,6 +40,7 @@ void Init_ext()
   init_backends(m);
   init_cuda(m);
   init_device(m);
+  init_generator(m, rb_cGenerator);
   init_ivalue(m, rb_cIValue);
   init_random(m);
 }

data/ext/torch/generator.cpp

@@ -0,0 +1,50 @@
+#include <torch/torch.h>
+
+#include <rice/rice.hpp>
+
+#include "utils.h"
+
+void init_generator(Rice::Module& m, Rice::Class& rb_cGenerator) {
+  // https://github.com/pytorch/pytorch/blob/master/torch/csrc/Generator.cpp
+  rb_cGenerator
+    .add_handler<torch::Error>(handle_error)
+    .define_singleton_function(
+      "new",
+      []() {
+        // TODO support more devices
+        return torch::make_generator<torch::CPUGeneratorImpl>();
+      })
+    .define_method(
+      "device",
+      [](torch::Generator& self) {
+        return self.device();
+      })
+    .define_method(
+      "initial_seed",
+      [](torch::Generator& self) {
+        return self.current_seed();
+      })
+    .define_method(
+      "manual_seed",
+      [](torch::Generator& self, uint64_t seed) {
+        self.set_current_seed(seed);
+        return self;
+      })
+    .define_method(
+      "seed",
+      [](torch::Generator& self) {
+        return self.seed();
+      })
+    .define_method(
+      "state",
+      [](torch::Generator& self) {
+        return self.get_state();
+      })
+    .define_method(
+      "state=",
+      [](torch::Generator& self, const torch::Tensor& state) {
+        self.set_state(state);
+      });
+
+  THPGeneratorClass = rb_cGenerator.value();
+}

data/ext/torch/ruby_arg_parser.cpp

@@ -2,6 +2,7 @@
 
 #include "ruby_arg_parser.h"
 
+VALUE THPGeneratorClass = Qnil;
 VALUE THPVariableClass = Qnil;
 
 static std::unordered_map<std::string, ParameterType> type_map = {
@@ -11,6 +12,7 @@ static std::unordered_map<std::string, ParameterType> type_map = {
   {"double", ParameterType::DOUBLE},
   {"complex", ParameterType::COMPLEX},
   {"TensorList", ParameterType::TENSOR_LIST},
+  {"c10::List<c10::optional<Tensor>>", ParameterType::TENSOR_LIST},
   {"IntArrayRef", ParameterType::INT_LIST},
   {"ArrayRef<double>", ParameterType::FLOAT_LIST},
   {"Generator", ParameterType::GENERATOR},
@@ -22,9 +24,14 @@ static std::unordered_map<std::string, ParameterType> type_map = {
   {"MemoryFormat", ParameterType::MEMORY_FORMAT},
   {"QScheme", ParameterType::QSCHEME},
   {"Device", ParameterType::DEVICE},
+  {"Stream", ParameterType::STREAM},
   {"std::string", ParameterType::STRING},
+  {"c10::string_view", ParameterType::STRING},
+  {"SymInt", ParameterType::SYM_INT},
   {"Dimname", ParameterType::DIMNAME},
+  {"SymIntArrayRef", ParameterType::SYM_INT_LIST},
   {"DimnameList", ParameterType::DIMNAME_LIST},
+  {"ScalarList", ParameterType::SCALAR_LIST}
 };
 
 static const std::unordered_map<std::string, std::vector<std::string>> numpy_compatibility_arg_names = {
@@ -116,6 +123,72 @@ bool is_tensor_list(VALUE obj, int argnum, bool throw_error) {
   return true;
 }
 
+static bool is_int_list(VALUE obj, int broadcast_size) {
+  if (RB_TYPE_P(obj, T_ARRAY)) {
+    auto len = RARRAY_LEN(obj);
+    if (len == 0) {
+      return true;
+    }
+
+    auto item = rb_ary_entry(obj, 0);
+    bool int_first = false;
+    if (THPUtils_checkIndex(item)) {
+      // we still have to check that the rest of items are NOT symint nodes
+      int_first = true;
+    }
+
+    // Make sure none of the later arguments are SymInt
+    // NB: do NOT check that the later arguments are ints, as this is
+    // BC-breaking for FX
+    // for (int i = 1; i < len; i++) {
+    //   if (torch::is_symint_node(
+    //           py::reinterpret_steal<py::object>(PySequence_GetItem(obj, i)))) {
+    //     return false;
+    //   }
+    // }
+
+    if (int_first) {
+      return true;
+    }
+
+    // NOTE: JIT tracer allows arbitrary scalar tensors to act as ints
+    // in an intlist argument. Even float or complex scalar tensors.
+    // return (
+    //     jit::tracer::isTracing() && THPVariable_Check(item.ptr()) &&
+    //     THPVariable_Unpack(item.ptr()).sizes() == c10::IntArrayRef{});
+    return false;
+  }
+  // if a size is specified (e.g. IntArrayRef[2]) we also allow passing a single
+  // int
+  return broadcast_size > 0 && THPUtils_checkLong(obj);
+}
+
+static bool is_int_or_symint(VALUE obj) {
+  return THPUtils_checkIndex(obj);
+}
+
+static bool is_int_or_symint_list(VALUE obj, int broadcast_size) {
+  if (RB_TYPE_P(obj, T_ARRAY)) {
+    if (RARRAY_LEN(obj) == 0) {
+      return true;
+    }
+    auto item = rb_ary_entry(obj, 0);
+
+    if (is_int_or_symint(item)) {
+      return true;
+    }
+    // NOTE: JIT tracer allows arbitrary scalar tensors to act as ints
+    // in an intlist argument. Even float or complex scalar tensors.
+    // return (
+    //     jit::tracer::isTracing() && THPVariable_Check(item.ptr()) &&
+    //     THPVariable_Unpack(item.ptr()).sizes() == c10::IntArrayRef{});
+    return false;
+  }
+  // if a size is specified (e.g. IntArrayRef[2]) we also allow passing a single
+  // int
+  return broadcast_size > 0 && THPUtils_checkLong(obj);
+}
+
 // argnum is needed for raising the TypeError, it's used in the error message.
 auto FunctionParameter::check(VALUE obj, int argnum) -> bool
 {
@@ -172,7 +245,7 @@ auto FunctionParameter::check(VALUE obj, int argnum) -> bool
       return size > 0 && FIXNUM_P(obj);
     }
     case ParameterType::FLOAT_LIST: return (RB_TYPE_P(obj, T_ARRAY));
-    case ParameterType::GENERATOR: return false; // return THPGenerator_Check(obj);
+    case ParameterType::GENERATOR: return THPGenerator_Check(obj);
     case ParameterType::BOOL: return obj == Qtrue || obj == Qfalse;
     case ParameterType::STORAGE: return false; // return isStorage(obj);
     // case ParameterType::PYOBJECT: return true;
@@ -182,6 +255,8 @@ auto FunctionParameter::check(VALUE obj, int argnum) -> bool
     case ParameterType::QSCHEME: return false; // return THPQScheme_Check(obj);
     case ParameterType::DEVICE: return RB_TYPE_P(obj, T_STRING); // TODO check device
     case ParameterType::STRING: return RB_TYPE_P(obj, T_STRING);
+    case ParameterType::SYM_INT: return is_int_or_symint(obj);
+    case ParameterType::SYM_INT_LIST: return is_int_or_symint_list(obj, size);
     default: throw std::runtime_error("unknown parameter type");
   }
 }
@@ -191,6 +266,7 @@ std::string FunctionParameter::type_name() const {
     case ParameterType::TENSOR: return "Tensor";
     case ParameterType::SCALAR: return "Number";
     case ParameterType::INT64: return "int";
+    case ParameterType::SYM_INT: return "SymInt";
     case ParameterType::DOUBLE: return "float";
     case ParameterType::COMPLEX: return "complex";
     case ParameterType::TENSOR_LIST: return "array of Tensors";
@@ -208,6 +284,8 @@ std::string FunctionParameter::type_name() const {
     case ParameterType::STRING: return "str";
     case ParameterType::DIMNAME: return "name";
     case ParameterType::DIMNAME_LIST: return "array of names";
+    case ParameterType::SCALAR_LIST: return "array of Scalars";
+    case ParameterType::SYM_INT_LIST: return "array of SymInts";
     default: throw std::runtime_error("unknown parameter type");
   }
 }
@@ -558,8 +636,14 @@ bool FunctionSignature::parse(VALUE self, VALUE args, VALUE kwargs, VALUE dst[],
 
   // if there is a single positional IntArrayRef argument, i.e. expand(..), view(...),
   // allow a var-args style IntArrayRef, so expand(5,3) behaves as expand((5,3))
-  if (max_pos_args == 1 && params[0].type_ == ParameterType::INT_LIST) {
+  int int_list_overload = false;
+  if (max_pos_args == 1 &&
+      (params[0].type_ == ParameterType::INT_LIST ||
+       params[0].type_ == ParameterType::SYM_INT_LIST)) {
     allow_varargs_intlist = true;
+    if (params[0].type_ == ParameterType::INT_LIST) {
+      int_list_overload = true;
+    }
   }
 
   if (nargs > max_pos_args && !allow_varargs_intlist) {
@@ -614,8 +698,10 @@ bool FunctionSignature::parse(VALUE self, VALUE args, VALUE kwargs, VALUE dst[],
     // XXX: the Variable check is necessary because sizes become tensors when
     // tracer is enabled. This behavior easily leads to ambiguities, and we
     // should avoid having complex signatures that make use of it...
-    } else if (allow_varargs_intlist && arg_pos == 0 && !is_kwd &&
-               THPUtils_checkIndex(obj)) {
+    } else if (
+        allow_varargs_intlist && arg_pos == 0 && !is_kwd &&
+        ((int_list_overload ? is_int_list(args, param.size)
+                            : is_int_or_symint_list(args, param.size)))) {
       // take all positional arguments as this parameter
       // e.g. permute(1, 2, 3) -> permute((1, 2, 3))
       dst[i++] = args;

data/ext/torch/ruby_arg_parser.h

@@ -11,9 +11,9 @@
 #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
+  TENSOR, SCALAR, INT64, SYM_INT, DOUBLE, COMPLEX, TENSOR_LIST, INT_LIST, GENERATOR,
+  BOOL, STORAGE, PYOBJECT, SCALARTYPE, LAYOUT, MEMORY_FORMAT, DEVICE, STREAM, STRING,
+  DIMNAME, DIMNAME_LIST, QSCHEME, FLOAT_LIST, SCALAR_LIST, SYM_INT_LIST
 };
 
 struct FunctionParameter {
@@ -84,7 +84,9 @@ struct RubyArgs {
   template<int N>
   inline std::array<at::Tensor, N> tensorlist_n(int i);
   inline std::vector<int64_t> intlist(int i);
+  inline std::vector<c10::SymInt> symintlist(int i);
   inline c10::OptionalArray<int64_t> intlistOptional(int i);
+  inline c10::OptionalArray<c10::SymInt> symintlistOptional(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);
@@ -93,6 +95,7 @@ struct RubyArgs {
   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<c10::SymInt> toSymIntOptional(int i);
   inline c10::optional<bool> toBoolOptional(int i);
   inline c10::optional<double> toDoubleOptional(int i);
   inline c10::OptionalArray<double> doublelistOptional(int i);
@@ -116,6 +119,7 @@ struct RubyArgs {
   inline c10::optional<c10::string_view> stringViewOptional(int i);
   // inline PyObject* pyobject(int i);
   inline int64_t toInt64(int i);
+  inline c10::SymInt toSymInt(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);
@@ -171,6 +175,19 @@ inline std::vector<int64_t> RubyArgs::intlist(int i) {
   return intlistWithDefault(i, signature.params[i].default_intlist);
 }
 
+inline std::vector<c10::SymInt> RubyArgs::symintlist(int i) {
+  if (NIL_P(args[i])) {
+    return c10::fmap(signature.params[i].default_intlist, [](int64_t di) {
+      return c10::SymInt(di);
+    });
+  }
+
+  // TODO improve
+  return c10::fmap(intlist(i), [](int64_t di) {
+    return c10::SymInt(di);
+  });
+}
+
 inline std::vector<int64_t> RubyArgs::intlistWithDefault(int i, std::vector<int64_t> default_intlist) {
   if (NIL_P(args[i])) return default_intlist;
   VALUE arg = args[i];
@@ -199,9 +216,14 @@ inline c10::OptionalArray<int64_t> RubyArgs::intlistOptional(int i) {
   return intlist(i);
 }
 
+inline c10::OptionalArray<c10::SymInt> RubyArgs::symintlistOptional(int i) {
+  if (NIL_P(args[i])) return {};
+  return symintlist(i);
+}
+
 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");
+  return Rice::detail::From_Ruby<torch::Generator>().convert(args[i]);
 }
 
 inline at::Storage RubyArgs::storage(int i) {
@@ -270,6 +292,11 @@ inline c10::optional<int64_t> RubyArgs::toInt64Optional(int i) {
   return toInt64(i);
 }
 
+inline c10::optional<c10::SymInt> RubyArgs::toSymIntOptional(int i) {
+  if (NIL_P(args[i])) return c10::nullopt;
+  return toSymInt(i);
+}
+
 inline c10::optional<bool> RubyArgs::toBoolOptional(int i) {
   if (NIL_P(args[i])) return c10::nullopt;
   return toBool(i);
@@ -376,6 +403,15 @@ inline int64_t RubyArgs::toInt64(int i) {
   return Rice::detail::From_Ruby<int64_t>().convert(args[i]);
 }
 
+inline c10::SymInt RubyArgs::toSymInt(int i) {
+  if (NIL_P(args[i])) {
+    return c10::SymInt(signature.params[i].default_int);
+  }
+
+  // TODO improve
+  return c10::SymInt(toInt64(i));
+}
+
 inline double RubyArgs::toDouble(int i) {
   if (NIL_P(args[i])) return signature.params[i].default_double;
   return Rice::detail::From_Ruby<double>().convert(args[i]);

data/ext/torch/utils.h

@@ -6,7 +6,7 @@
 #include <rice/stl.hpp>
 
 static_assert(
-  TORCH_VERSION_MAJOR == 1 && TORCH_VERSION_MINOR == 12,
+  TORCH_VERSION_MAJOR == 1 && TORCH_VERSION_MINOR == 13,
   "Incompatible LibTorch version"
 );
 
@@ -17,6 +17,7 @@ inline void handle_error(torch::Error const & ex) {
 
 // keep THP prefix for now to make it easier to compare code
 
+extern VALUE THPGeneratorClass;
 extern VALUE THPVariableClass;
 
 inline VALUE THPUtils_internSymbol(const std::string& str) {
@@ -36,10 +37,18 @@ inline bool THPUtils_checkIndex(VALUE obj) {
   return FIXNUM_P(obj);
 }
 
+inline bool THPUtils_checkLong(VALUE obj) {
+  return FIXNUM_P(obj);
+}
+
 inline bool THPUtils_checkScalar(VALUE obj) {
   return FIXNUM_P(obj) || RB_FLOAT_TYPE_P(obj) || RB_TYPE_P(obj, T_COMPLEX);
 }
 
+inline bool THPGenerator_Check(VALUE obj) {
+  return rb_obj_is_kind_of(obj, THPGeneratorClass);
+}
+
 inline bool THPVariable_Check(VALUE obj) {
   return rb_obj_is_kind_of(obj, THPVariableClass);
 }

data/lib/torch/version.rb

@@ -1,3 +1,3 @@
 module Torch
-  VERSION = "0.11.2"
+  VERSION = "0.12.1"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: torch-rb
 version: !ruby/object:Gem::Version
-  version: 0.11.2
+  version: 0.12.1
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2022-09-25 00:00:00.000000000 Z
+date: 2023-01-30 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rice
@@ -44,6 +44,7 @@ files:
 - ext/torch/extconf.rb
 - ext/torch/fft.cpp
 - ext/torch/fft_functions.h
+- ext/torch/generator.cpp
 - ext/torch/ivalue.cpp
 - ext/torch/linalg.cpp
 - ext/torch/linalg_functions.h
@@ -229,7 +230,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.3.7
+rubygems_version: 3.4.1
 signing_key:
 specification_version: 4
 summary: Deep learning for Ruby, powered by LibTorch