torch-rb 0.5.3 → 0.6.0

data/ext/torch/extconf.rb

@@ -11,7 +11,6 @@ apple_clang = RbConfig::CONFIG["CC_VERSION_MESSAGE"] =~ /apple clang/i
 
 # check omp first
 if have_library("omp") || have_library("gomp")
-  $CXXFLAGS += " -DAT_PARALLEL_OPENMP=1"
   $CXXFLAGS += " -Xclang" if apple_clang
   $CXXFLAGS += " -fopenmp"
 end

data/ext/torch/ivalue.cpp

@@ -0,0 +1,134 @@
+#include <torch/torch.h>
+
+#include <rice/Array.hpp>
+#include <rice/Constructor.hpp>
+#include <rice/Hash.hpp>
+#include <rice/Module.hpp>
+#include <rice/String.hpp>
+
+#include "utils.h"
+
+void init_ivalue(Rice::Module& m) {
+  // https://pytorch.org/cppdocs/api/structc10_1_1_i_value.html
+  Rice::define_class_under<torch::IValue>(m, "IValue")
+    .add_handler<torch::Error>(handle_error)
+    .define_constructor(Rice::Constructor<torch::IValue>())
+    .define_method("bool?", &torch::IValue::isBool)
+    .define_method("bool_list?", &torch::IValue::isBoolList)
+    .define_method("capsule?", &torch::IValue::isCapsule)
+    .define_method("custom_class?", &torch::IValue::isCustomClass)
+    .define_method("device?", &torch::IValue::isDevice)
+    .define_method("double?", &torch::IValue::isDouble)
+    .define_method("double_list?", &torch::IValue::isDoubleList)
+    .define_method("future?", &torch::IValue::isFuture)
+    // .define_method("generator?", &torch::IValue::isGenerator)
+    .define_method("generic_dict?", &torch::IValue::isGenericDict)
+    .define_method("list?", &torch::IValue::isList)
+    .define_method("int?", &torch::IValue::isInt)
+    .define_method("int_list?", &torch::IValue::isIntList)
+    .define_method("module?", &torch::IValue::isModule)
+    .define_method("none?", &torch::IValue::isNone)
+    .define_method("object?", &torch::IValue::isObject)
+    .define_method("ptr_type?", &torch::IValue::isPtrType)
+    .define_method("py_object?", &torch::IValue::isPyObject)
+    .define_method("r_ref?", &torch::IValue::isRRef)
+    .define_method("scalar?", &torch::IValue::isScalar)
+    .define_method("string?", &torch::IValue::isString)
+    .define_method("tensor?", &torch::IValue::isTensor)
+    .define_method("tensor_list?", &torch::IValue::isTensorList)
+    .define_method("tuple?", &torch::IValue::isTuple)
+    .define_method(
+      "to_bool",
+      *[](torch::IValue& self) {
+        return self.toBool();
+      })
+    .define_method(
+      "to_double",
+      *[](torch::IValue& self) {
+        return self.toDouble();
+      })
+    .define_method(
+      "to_int",
+      *[](torch::IValue& self) {
+        return self.toInt();
+      })
+    .define_method(
+      "to_list",
+      *[](torch::IValue& self) {
+        auto list = self.toListRef();
+        Rice::Array obj;
+        for (auto& elem : list) {
+          obj.push(to_ruby<torch::IValue>(torch::IValue{elem}));
+        }
+        return obj;
+      })
+    .define_method(
+      "to_string_ref",
+      *[](torch::IValue& self) {
+        return self.toStringRef();
+      })
+    .define_method(
+      "to_tensor",
+      *[](torch::IValue& self) {
+        return self.toTensor();
+      })
+    .define_method(
+      "to_generic_dict",
+      *[](torch::IValue& self) {
+        auto dict = self.toGenericDict();
+        Rice::Hash obj;
+        for (auto& pair : dict) {
+          obj[to_ruby<torch::IValue>(torch::IValue{pair.key()})] = to_ruby<torch::IValue>(torch::IValue{pair.value()});
+        }
+        return obj;
+      })
+    .define_singleton_method(
+      "from_tensor",
+      *[](torch::Tensor& v) {
+        return torch::IValue(v);
+      })
+    // TODO create specialized list types?
+    .define_singleton_method(
+      "from_list",
+      *[](Rice::Array obj) {
+        c10::impl::GenericList list(c10::AnyType::get());
+        for (auto entry : obj) {
+          list.push_back(from_ruby<torch::IValue>(entry));
+        }
+        return torch::IValue(list);
+      })
+    .define_singleton_method(
+      "from_string",
+      *[](Rice::String v) {
+        return torch::IValue(v.str());
+      })
+    .define_singleton_method(
+      "from_int",
+      *[](int64_t v) {
+        return torch::IValue(v);
+      })
+    .define_singleton_method(
+      "from_double",
+      *[](double v) {
+        return torch::IValue(v);
+      })
+    .define_singleton_method(
+      "from_bool",
+      *[](bool v) {
+        return torch::IValue(v);
+      })
+    // see https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/python/pybind_utils.h
+    // createGenericDict and toIValue
+    .define_singleton_method(
+      "from_dict",
+      *[](Rice::Hash obj) {
+        auto key_type = c10::AnyType::get();
+        auto value_type = c10::AnyType::get();
+        c10::impl::GenericDict elems(key_type, value_type);
+        elems.reserve(obj.size());
+        for (auto entry : obj) {
+          elems.insert(from_ruby<torch::IValue>(entry.first), from_ruby<torch::IValue>((Rice::Object) entry.second));
+        }
+        return torch::IValue(std::move(elems));
+      });
+}

data/ext/torch/nn.cpp

@@ -0,0 +1,114 @@
+#include <torch/torch.h>
+
+#include <rice/Module.hpp>
+
+#include "nn_functions.h"
+#include "templates.h"
+#include "utils.h"
+
+// need to make a distinction between parameters and tensors
+class Parameter: public torch::autograd::Variable {
+  public:
+    Parameter(Tensor&& t) : torch::autograd::Variable(t) { }
+};
+
+void init_nn(Rice::Module& m) {
+  auto rb_mNN = Rice::define_module_under(m, "NN");
+  rb_mNN.add_handler<torch::Error>(handle_error);
+  add_nn_functions(rb_mNN);
+
+  Rice::define_module_under(rb_mNN, "Init")
+    .add_handler<torch::Error>(handle_error)
+    .define_singleton_method(
+      "_calculate_gain",
+      *[](NonlinearityType nonlinearity, double param) {
+        return torch::nn::init::calculate_gain(nonlinearity, param);
+      })
+    .define_singleton_method(
+      "_uniform!",
+      *[](Tensor tensor, double low, double high) {
+        return torch::nn::init::uniform_(tensor, low, high);
+      })
+    .define_singleton_method(
+      "_normal!",
+      *[](Tensor tensor, double mean, double std) {
+        return torch::nn::init::normal_(tensor, mean, std);
+      })
+    .define_singleton_method(
+      "_constant!",
+      *[](Tensor tensor, Scalar value) {
+        return torch::nn::init::constant_(tensor, value);
+      })
+    .define_singleton_method(
+      "_ones!",
+      *[](Tensor tensor) {
+        return torch::nn::init::ones_(tensor);
+      })
+    .define_singleton_method(
+      "_zeros!",
+      *[](Tensor tensor) {
+        return torch::nn::init::zeros_(tensor);
+      })
+    .define_singleton_method(
+      "_eye!",
+      *[](Tensor tensor) {
+        return torch::nn::init::eye_(tensor);
+      })
+    .define_singleton_method(
+      "_dirac!",
+      *[](Tensor tensor) {
+        return torch::nn::init::dirac_(tensor);
+      })
+    .define_singleton_method(
+      "_xavier_uniform!",
+      *[](Tensor tensor, double gain) {
+        return torch::nn::init::xavier_uniform_(tensor, gain);
+      })
+    .define_singleton_method(
+      "_xavier_normal!",
+      *[](Tensor tensor, double gain) {
+        return torch::nn::init::xavier_normal_(tensor, gain);
+      })
+    .define_singleton_method(
+      "_kaiming_uniform!",
+      *[](Tensor tensor, double a, FanModeType mode, NonlinearityType nonlinearity) {
+        return torch::nn::init::kaiming_uniform_(tensor, a, mode, nonlinearity);
+      })
+    .define_singleton_method(
+      "_kaiming_normal!",
+      *[](Tensor tensor, double a, FanModeType mode, NonlinearityType nonlinearity) {
+        return torch::nn::init::kaiming_normal_(tensor, a, mode, nonlinearity);
+      })
+    .define_singleton_method(
+      "_orthogonal!",
+      *[](Tensor tensor, double gain) {
+        return torch::nn::init::orthogonal_(tensor, gain);
+      })
+    .define_singleton_method(
+      "_sparse!",
+      *[](Tensor tensor, double sparsity, double std) {
+        return torch::nn::init::sparse_(tensor, sparsity, std);
+      });
+
+  Rice::define_class_under<Parameter, torch::Tensor>(rb_mNN, "Parameter")
+    .add_handler<torch::Error>(handle_error)
+    .define_method(
+      "grad",
+      *[](Parameter& self) {
+        auto grad = self.grad();
+        return grad.defined() ? to_ruby<torch::Tensor>(grad) : Nil;
+      })
+    .define_method(
+      "grad=",
+      *[](Parameter& self, torch::Tensor& grad) {
+        self.mutable_grad() = grad;
+      })
+    .define_singleton_method(
+      "_make_subclass",
+      *[](Tensor& rd, bool requires_grad) {
+        auto data = rd.detach();
+        data.unsafeGetTensorImpl()->set_allow_tensor_metadata_change(true);
+        auto var = data.set_requires_grad(requires_grad);
+        return Parameter(std::move(var));
+      });
+}

data/ext/torch/nn_functions.h

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

data/ext/torch/random.cpp

@@ -0,0 +1,22 @@
+#include <torch/torch.h>
+
+#include <rice/Module.hpp>
+
+#include "utils.h"
+
+void init_random(Rice::Module& m) {
+  Rice::define_module_under(m, "Random")
+    .add_handler<torch::Error>(handle_error)
+    .define_singleton_method(
+      "initial_seed",
+      *[]() {
+        return at::detail::getDefaultCPUGenerator().current_seed();
+      })
+    .define_singleton_method(
+      "seed",
+      *[]() {
+        // TODO set for CUDA when available
+        auto generator = at::detail::getDefaultCPUGenerator();
+        return generator.seed();
+      });
+}

data/ext/torch/ruby_arg_parser.cpp

@@ -487,7 +487,7 @@ static void extra_kwargs(FunctionSignature& signature, VALUE kwargs, ssize_t num
 
 VALUE missing = Qundef;
 
-bool FunctionSignature::parse(VALUE self, VALUE args, VALUE kwargs, std::vector<VALUE> &dst,  // NOLINT
+bool FunctionSignature::parse(VALUE self, VALUE args, VALUE kwargs, VALUE dst[],  // NOLINT
                               bool raise_exception) {
   auto nargs = NIL_P(args) ? 0 : RARRAY_LEN(args);
   ssize_t remaining_kwargs = NIL_P(kwargs) ? 0 :  RHASH_SIZE(kwargs);

data/ext/torch/ruby_arg_parser.h

@@ -2,6 +2,8 @@
 
 #pragma once
 
+#include <sstream>
+
 #include <torch/torch.h>
 #include <rice/Exception.hpp>
 
@@ -46,7 +48,7 @@ struct FunctionParameter {
 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);
+  bool parse(VALUE self, VALUE args, VALUE kwargs, VALUE dst[], bool raise_exception);
 
   std::string toString() const;
 
@@ -63,13 +65,13 @@ struct FunctionSignature {
 };
 
 struct RubyArgs {
-  RubyArgs(const FunctionSignature& signature, std::vector<VALUE> &args)
+  RubyArgs(const FunctionSignature& signature, VALUE* args)
     : signature(signature)
     , args(args)
     , idx(signature.index) {}
 
   const FunctionSignature& signature;
-  std::vector<VALUE> args;
+  VALUE* args;
   int idx;
 
   inline at::Tensor tensor(int i);
@@ -328,6 +330,12 @@ inline bool RubyArgs::isNone(int i) {
   return NIL_P(args[i]);
 }
 
+template<int N>
+struct ParsedArgs {
+  ParsedArgs() : args() { }
+  VALUE args[N];
+};
+
 struct RubyArgParser {
   std::vector<FunctionSignature> signatures_;
   std::string function_name;
@@ -356,7 +364,15 @@ struct RubyArgParser {
         });
     }
 
-    RubyArgs parse(VALUE self, int argc, VALUE* argv, std::vector<VALUE> &parsed_args) {
+    template<int N>
+    inline RubyArgs parse(VALUE self, int argc, VALUE* argv, ParsedArgs<N> &dst) {
+      if (N < max_args) {
+        rb_raise(rb_eArgError, "RubyArgParser: dst ParsedArgs buffer does not have enough capacity, expected %d (got %d)", (int)max_args, N);
+      }
+      return raw_parse(self, argc, argv, dst.args);
+    }
+
+    inline RubyArgs raw_parse(VALUE self, int argc, VALUE* argv, VALUE parsed_args[]) {
       VALUE args, kwargs;
       rb_scan_args(argc, argv, "*:", &args, &kwargs);
 
@@ -378,7 +394,7 @@ struct RubyArgParser {
       rb_raise(rb_eArgError, "No matching signatures");
     }
 
-    void print_error(VALUE self, VALUE args, VALUE kwargs, std::vector<VALUE>& parsed_args) {
+    void print_error(VALUE self, VALUE args, VALUE kwargs, 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;

data/ext/torch/templates.h

@@ -44,7 +44,7 @@ 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++) {
+  for (long i = 0; i < a.size(); i++) {
     vec[i] = from_ruby<int64_t>(a[i]);
   }
   return vec;
@@ -56,7 +56,7 @@ std::vector<Tensor> from_ruby<std::vector<Tensor>>(Object x)
 {
   Array a = Array(x);
   std::vector<Tensor> vec(a.size());
-  for (size_t i = 0; i < a.size(); i++) {
+  for (long i = 0; i < a.size(); i++) {
     vec[i] = from_ruby<Tensor>(a[i]);
   }
   return vec;

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);
+      });
+}