torch-rb 0.11.2 → 0.12.0

data/ext/torch/ruby_arg_parser.cpp

@@ -11,6 +11,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 +23,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 +122,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
 {
@@ -182,6 +254,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 +265,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 +283,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 +635,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 +697,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,6 +216,11 @@ 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");
@@ -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"
 );
 
@@ -36,6 +36,10 @@ 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);
 }

data/lib/torch/version.rb

@@ -1,3 +1,3 @@
 module Torch
-  VERSION = "0.11.2"
+  VERSION = "0.12.0"
 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.0
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2022-09-25 00:00:00.000000000 Z
+date: 2022-11-05 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rice