torch-rb 0.1.3 → 0.1.8

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e7f715179c9a84dc7399b80d93fd61f2bbb58a0156e6084dc4abb23e1d4a1b52
-  data.tar.gz: 6928379ae7c92a77ad9dde4f4224ec33c6f8575a9b77585c0147e4f5361021de
+  metadata.gz: fca87cb9b6d255287e9fafadf786c113798abbe76b36c82b8271b79cfbf3c2b9
+  data.tar.gz: 4813c71f5ad6d078e78da03cf59f8036e9e76258ffb67f538899bba146dcba2a
 SHA512:
-  metadata.gz: 9911a9e86d93f1e410776c44fdb3cd9aa06c83d1f0e42fdab8530970bea6520aed7906e96fb8243efd6b957453ebc13678b2b92e4c85b54407030a32c6196e08
-  data.tar.gz: 0d080f5458a5dcf8fee19ce5e2e342bf6269432de6e78d923036232963ebb80daeea993c0bbf4af2d6da46593ac28a72a8232020a9fcb48acc3276c9e1ebebf3
+  metadata.gz: 22c7150e6a7d9132c40c67819beecc6b8c69b268bd227a8e4aa324ef5e2707004691d5b65dcd4ba1ac537bfaf783947da7e5a323417cffcbf7d348768c40b7c6
+  data.tar.gz: 8a86c6b68efe6ad85a261d7033b87f040c22b2c670a0238accd6246274caed17b86d7b424441bba80c5ea67ec1bf53b05444dfb0c45ea5b8a52806d0ce19ec1e

data/CHANGELOG.md

@@ -1,3 +1,33 @@
+## 0.1.8 (2020-01-17)
+
+- Added support for libtorch 1.4.0
+- Dropped support for libtorch 1.3.1
+
+## 0.1.7 (2020-01-10)
+
+- Fixed installation error with Ruby 2.7
+
+## 0.1.6 (2019-12-09)
+
+- Added recurrent layers
+- Added more pooling layers
+- Added normalization layers
+
+## 0.1.5 (2019-12-06)
+
+- Added many more functions
+- Added tensor classes - `FloatTensor`, `LongTensor`, etc
+- Improved modules
+
+## 0.1.4 (2019-12-01)
+
+- Added distance functions
+- Added more activations
+- Added more linear layers
+- Added more loss functions
+- Added more init methods
+- Added support for tensor assignment
+
 ## 0.1.3 (2019-11-30)
 
 - Changed to BSD 3-Clause license to match PyTorch

data/README.md

@@ -20,6 +20,8 @@ Add this line to your application’s Gemfile:
 gem 'torch-rb'
 ```
 
+It can take a few minutes to compile the extension.
+
 ## Getting Started
 
 This library follows the [PyTorch API](https://pytorch.org/docs/stable/torch.html). There are a few changes to make it more Ruby-like:
@@ -28,7 +30,7 @@ This library follows the [PyTorch API](https://pytorch.org/docs/stable/torch.htm
 - Methods that return booleans use `?` instead of `is_`  (`tensor?` instead of `is_tensor`)
 - Numo is used instead of NumPy (`x.numo` instead of `x.numpy()`)
 
-Many methods and options are missing at the moment. PRs welcome!
+Some methods and options are missing at the moment. PRs welcome!
 
 ## Tutorial
 
@@ -365,8 +367,9 @@ Here’s a list of functions to create tensors (descriptions from the [C++ docs]
 
 Here are a few full examples:
 
-- [Image classification with MNIST](examples/mnist)
+- [Image classification with MNIST](examples/mnist) ([日本語版](https://qiita.com/kojix2/items/c19c36dc1bf73ea93409))
 - [Collaborative filtering with MovieLens](examples/movielens)
+- [Sequence models and word embeddings](examples/nlp)
 
 ## LibTorch Installation
 

data/ext/torch/ext.cpp

@@ -6,137 +6,35 @@
 #include <rice/Class.hpp>
 #include <rice/Constructor.hpp>
 
-using namespace Rice;
-
-template<>
-inline
-long long from_ruby<long long>(Object x)
-{
-  return NUM2LL(x);
-}
-
-template<>
-inline
-Object to_ruby<long long>(long long const & x)
-{
-  return LL2NUM(x);
-}
-
-template<>
-inline
-unsigned long long from_ruby<unsigned long long>(Object x)
-{
-  return NUM2ULL(x);
-}
-
-template<>
-inline
-Object to_ruby<unsigned long long>(unsigned long long const & x)
-{
-  return ULL2NUM(x);
-}
-
-template<>
-inline
-short from_ruby<short>(Object x)
-{
-  return NUM2SHORT(x);
-}
-
-template<>
-inline
-Object to_ruby<short>(short const & x)
-{
-  return INT2NUM(x);
-}
-
-template<>
-inline
-unsigned short from_ruby<unsigned short>(Object x)
-{
-  return NUM2USHORT(x);
-}
+#include "templates.hpp"
 
-template<>
-inline
-Object to_ruby<unsigned short>(unsigned short const & x)
-{
-  return UINT2NUM(x);
-}
+// generated with:
+// rake generate:functions
+#include "torch_functions.hpp"
+#include "tensor_functions.hpp"
+#include "nn_functions.hpp"
 
-// 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);
-}
+using namespace Rice;
 
-// for now
-class Scalar {
-  torch::Scalar value;
+// need to make a distinction between parameters and tensors
+class Parameter: public torch::autograd::Variable {
   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;
-    }
+    Parameter(Tensor&& t) : torch::autograd::Variable(t) { }
 };
 
-template<>
-inline
-Scalar from_ruby<Scalar>(Object x)
+extern "C"
+void Init_ext()
 {
-  return Scalar(x);
-}
+  Module rb_mTorch = define_module("Torch");
+  add_torch_functions(rb_mTorch);
 
-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);
-    }
-};
+  Class rb_cTensor = define_class_under<torch::Tensor>(rb_mTorch, "Tensor");
+  add_tensor_functions(rb_cTensor);
 
-template<>
-inline
-TensorList from_ruby<TensorList>(Object x)
-{
-  return TensorList(x);
-}
+  Module rb_mNN = define_module_under(rb_mTorch, "NN");
+  add_nn_functions(rb_mNN);
 
-extern "C"
-void Init_ext()
-{
-  Module rb_mTorch = define_module("Torch")
-    .define_singleton_method(
+  rb_mTorch.define_singleton_method(
       "grad_enabled?",
       *[]() {
         return torch::GradMode::is_enabled();
@@ -146,11 +44,6 @@ void Init_ext()
       *[](bool enabled) {
         torch::GradMode::set_enabled(enabled);
       })
-    .define_singleton_method(
-      "floating_point?",
-      *[](torch::Tensor& input) {
-        return torch::is_floating_point(input);
-      })
     .define_singleton_method(
       "manual_seed",
       *[](uint64_t seed) {
@@ -219,277 +112,17 @@ void Init_ext()
       })
     // begin operations
     .define_singleton_method(
-      "_mean",
-      *[](torch::Tensor& input) {
-        return torch::mean(input);
-      })
-    .define_singleton_method(
-      "_mean_dim",
-      *[](torch::Tensor& input, int64_t dim, bool keepdim) {
-        return torch::mean(input, dim, keepdim);
-      })
-    .define_singleton_method(
-      "_sum",
-      *[](torch::Tensor& input) {
-        return torch::sum(input);
-      })
-    .define_singleton_method(
-      "_sum_dim",
-      *[](torch::Tensor& input, int64_t dim, bool keepdim) {
-        return torch::sum(input, dim, keepdim);
-      })
-    .define_singleton_method(
-      "_argmax",
-      *[](torch::Tensor& input) {
-        return torch::argmax(input);
-      })
-    .define_singleton_method(
-      "_argmax_dim",
-      *[](torch::Tensor& input, int64_t dim, bool keepdim) {
-        return torch::argmax(input, dim, keepdim);
-      })
-    .define_singleton_method(
-      "_cat",
-      *[](TensorList tensors, int64_t dim) {
-        return torch::cat(tensors, dim);
-      })
-    .define_singleton_method(
-      "_norm",
-      *[](torch::Tensor& input) {
-        return torch::norm(input);
-      })
-    .define_singleton_method(
-      "_min",
-      *[](torch::Tensor& input) {
-        return torch::min(input);
-      })
-    .define_singleton_method(
-      "_max",
-      *[](torch::Tensor& input) {
-        return torch::max(input);
-      })
-    .define_singleton_method(
-      "_max_out",
-      *[](torch::Tensor &max, torch::Tensor &max_indices, const torch::Tensor &input, int64_t dim, bool keepdim) {
-        // TODO add return value
-        torch::_max_out(max, max_indices, input, dim, keepdim);
-      })
-    .define_singleton_method(
-      "_sqrt",
-      *[](torch::Tensor& input) {
-        return torch::sqrt(input);
-      })
-    .define_singleton_method(
-      "_exp",
-      *[](torch::Tensor& input) {
-        return torch::exp(input);
-      })
-    .define_singleton_method(
-      "_log",
-      *[](torch::Tensor& input) {
-        return torch::log(input);
-      })
-    .define_singleton_method(
-      "_sign",
-      *[](torch::Tensor& input) {
-        return torch::sign(input);
-      })
-    .define_singleton_method(
-      "_unsqueeze",
-      *[](torch::Tensor& input, int64_t dim) {
-        return torch::unsqueeze(input, dim);
-      })
-    .define_singleton_method(
-      "_dot",
-      *[](torch::Tensor& input, torch::Tensor& tensor) {
-        return torch::dot(input, tensor);
-      })
-    .define_singleton_method(
-      "_matmul",
-      *[](torch::Tensor& input, torch::Tensor& other) {
-        return torch::matmul(input, other);
-      })
-    .define_singleton_method(
-      "_eq",
-      *[](torch::Tensor& input, torch::Tensor& other) {
-        return torch::eq(input, other);
-      })
-    .define_singleton_method(
-      "_gt",
-      // TODO support tensors
-      *[](torch::Tensor& input, Scalar other) {
-        return torch::gt(input, other);
-      })
-    .define_singleton_method(
-      "_lt",
-      // TODO support tensors
-      *[](torch::Tensor& input, Scalar other) {
-        return torch::lt(input, other);
-      })
-    .define_singleton_method(
-      "_add",
-      *[](torch::Tensor& input, torch::Tensor& other) {
-        return torch::add(input, other);
-      })
-    .define_singleton_method(
-      "_add_scalar",
-      *[](torch::Tensor& input, Scalar other) {
-        return torch::add(input, other);
-      })
-    .define_singleton_method(
-      "_add_out",
-      *[](torch::Tensor& out, torch::Tensor& input, torch::Tensor& other) {
-        return torch::add_out(out, input, other);
-      })
-    .define_singleton_method(
-      "_sub",
-      *[](torch::Tensor& input, torch::Tensor& other) {
-        return torch::sub(input, other);
-      })
-    .define_singleton_method(
-      "_sub_scalar",
-      *[](torch::Tensor& input, Scalar other) {
-        return torch::sub(input, other);
-      })
-    .define_singleton_method(
-      "_mul",
-      *[](torch::Tensor& input, torch::Tensor& other) {
-        return torch::mul(input, other);
-      })
-    .define_singleton_method(
-      "_mul_scalar",
-      *[](torch::Tensor& input, Scalar other) {
-        return torch::mul(input, other);
-      })
-    .define_singleton_method(
-      "_div",
-      *[](torch::Tensor& input, torch::Tensor& other) {
-        return torch::div(input, other);
-      })
-    .define_singleton_method(
-      "_div_scalar",
-      *[](torch::Tensor& input, Scalar other) {
-        return torch::div(input, other);
-      })
-    .define_singleton_method(
-      "_remainder",
-      *[](torch::Tensor& input, torch::Tensor& other) {
-        return torch::remainder(input, other);
-      })
-    .define_singleton_method(
-      "_remainder_scalar",
-      *[](torch::Tensor& input, Scalar other) {
-        return torch::remainder(input, other);
-      })
-    .define_singleton_method(
-      "_pow",
-      *[](torch::Tensor& input, Scalar exponent) {
-        return torch::pow(input, exponent);
-      })
-    .define_singleton_method(
-      "_abs",
-      *[](torch::Tensor& input) {
-        return torch::abs(input);
-      })
-    .define_singleton_method(
-      "_neg",
-      *[](torch::Tensor& input) {
-        return torch::neg(input);
+      "_save",
+      *[](const Tensor &value) {
+        auto v = torch::pickle_save(value);
+        std::string str(v.begin(), v.end());
+        return str;
       })
     .define_singleton_method(
-      "_reshape",
-      *[](torch::Tensor& input, IntArrayRef shape) {
-        return torch::reshape(input, shape);
+      "_binary_cross_entropy_with_logits",
+      *[](const Tensor &input, const Tensor &target, OptionalTensor weight, OptionalTensor pos_weight, MyReduction reduction) {
+        return torch::binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction);
       })
-    .define_singleton_method(
-      "_flatten",
-      *[](torch::Tensor& input, int64_t start_dim, int64_t end_dim) {
-        return torch::flatten(input, start_dim, end_dim);
-      })
-    .define_singleton_method(
-      "relu",
-      *[](torch::Tensor& input) {
-        return torch::relu(input);
-      })
-    .define_singleton_method(
-      "conv2d",
-      *[](torch::Tensor& input, torch::Tensor& weight, torch::Tensor& bias, IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, int64_t groups) {
-        return torch::conv2d(input, weight, bias, stride, padding, dilation, groups);
-      })
-    .define_singleton_method(
-      "linear",
-      *[](torch::Tensor& input, torch::Tensor& weight, torch::Tensor& bias) {
-        return torch::linear(input, weight, bias);
-      })
-    .define_singleton_method(
-      "max_pool2d",
-      *[](torch::Tensor& input, IntArrayRef kernel_size) {
-        return torch::max_pool2d(input, kernel_size);
-      })
-    .define_singleton_method(
-      "avg_pool2d",
-      *[](torch::Tensor& input, IntArrayRef kernel_size) {
-        return torch::avg_pool2d(input, kernel_size);
-      })
-    .define_singleton_method(
-      "_dropout",
-      *[](torch::Tensor& input, float p, bool train) {
-        return torch::dropout(input, p, train);
-      })
-    .define_singleton_method(
-      "_dropout!",
-      *[](torch::Tensor& input, float p, bool train) {
-        return torch::dropout_(input, p, train);
-      })
-    .define_singleton_method(
-      "_feature_dropout",
-      *[](torch::Tensor& input, float p, bool train) {
-        return torch::feature_dropout(input, p, train);
-      })
-    .define_singleton_method(
-      "_feature_dropout!",
-      *[](torch::Tensor& input, float p, bool train) {
-        return torch::feature_dropout_(input, p, train);
-      })
-    .define_singleton_method(
-      "_alpha_dropout",
-      *[](torch::Tensor& input, float p, bool train) {
-        return torch::alpha_dropout(input, p, train);
-      })
-    .define_singleton_method(
-      "_alpha_dropout!",
-      *[](torch::Tensor& input, float p, bool train) {
-        return torch::alpha_dropout_(input, p, train);
-      })
-    .define_singleton_method(
-      "_feature_alpha_dropout",
-      *[](torch::Tensor& input, float p, bool train) {
-        return torch::feature_alpha_dropout(input, p, train);
-      })
-    .define_singleton_method(
-      "_feature_alpha_dropout!",
-      *[](torch::Tensor& input, float p, bool train) {
-        return torch::feature_alpha_dropout_(input, p, train);
-      })
-    .define_singleton_method(
-      "_embedding",
-      // weight and indices are swapped from Python interface
-      *[](const torch::Tensor &indices, const torch::Tensor &weight, int64_t padding_idx, bool scale_grad_by_freq, bool sparse) {
-        return torch::embedding(weight, indices, padding_idx, scale_grad_by_freq, sparse);
-      })
-    .define_singleton_method(
-      "mse_loss",
-      *[](torch::Tensor& input, torch::Tensor& target, std::string reduction) {
-        auto red = reduction == "mean" ? Reduction::Mean : Reduction::Sum;
-        return torch::mse_loss(input, target, red);
-      })
-    .define_singleton_method(
-      "nll_loss",
-      *[](torch::Tensor& input, torch::Tensor& target, std::string reduction) {
-        auto red = reduction == "mean" ? Reduction::Mean : Reduction::Sum;
-        return torch::nll_loss(input, target, {}, red);
-      })
-    .define_singleton_method("numel", &torch::numel)
     .define_singleton_method(
       "_from_blob",
       *[](String s, IntArrayRef size, const torch::TensorOptions &options) {
@@ -500,187 +133,86 @@ void Init_ext()
       "_tensor",
       *[](Object o, IntArrayRef size, const torch::TensorOptions &options) {
         Array a = Array(o);
-        std::vector<float> vec;
-        for (size_t i = 0; i < a.size(); i++) {
-          vec.push_back(from_ruby<float>(a[i]));
+        auto dtype = options.dtype();
+        torch::Tensor t;
+        if (dtype == torch::kBool) {
+          throw std::runtime_error("Cannot create bool from tensor method yet");
+        } else {
+          std::vector<float> vec;
+          for (size_t i = 0; i < a.size(); i++) {
+            vec.push_back(from_ruby<float>(a[i]));
+          }
+          // hack for requires_grad error
+          if (options.requires_grad()) {
+            t = torch::tensor(vec, options.requires_grad(c10::nullopt));
+            t.set_requires_grad(true);
+          } else {
+            t = torch::tensor(vec, options);
+          }
         }
-        return torch::tensor(vec, options).reshape(size);
+        return t.reshape(size);
       });
 
-  Class rb_cTensor = define_class_under<torch::Tensor>(rb_mTorch, "Tensor")
+  rb_cTensor
     .define_method("cuda?", &torch::Tensor::is_cuda)
-    .define_method("distributed?", &torch::Tensor::is_distributed)
-    .define_method("complex?", &torch::Tensor::is_complex)
-    .define_method("floating_point?", &torch::Tensor::is_floating_point)
-    .define_method("signed?", &torch::Tensor::is_signed)
     .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("view_as", &torch::Tensor::view_as)
     .define_method(
       "addcmul!",
-      *[](torch::Tensor& self, Scalar value, const torch::Tensor & tensor1, const torch::Tensor & tensor2) {
+      *[](Tensor& self, Scalar value, const Tensor & tensor1, const Tensor & tensor2) {
         return self.addcmul_(tensor1, tensor2, value);
       })
     .define_method(
       "addcdiv!",
-      *[](torch::Tensor& self, Scalar value, const torch::Tensor & tensor1, const torch::Tensor & tensor2) {
+      *[](Tensor& self, Scalar value, const Tensor & tensor1, const Tensor & tensor2) {
         return self.addcdiv_(tensor1, tensor2, value);
       })
-    .define_method(
-      "zero!",
-      *[](torch::Tensor& self) {
-        return self.zero_();
-      })
-    .define_method(
-      "detach!",
-      *[](torch::Tensor& self) {
-        return self.detach_();
-      })
-    .define_method(
-      "_select",
-      *[](torch::Tensor& self, int64_t dim, int64_t index) {
-        return self.select(dim, index);
-      })
-    .define_method(
-      "_slice",
-      *[](torch::Tensor& self, int64_t dim, int64_t start, int64_t end, int64_t step) {
-        return self.slice(dim, start, end, step);
-      })
     .define_method(
       "_requires_grad!",
-      *[](torch::Tensor& self, bool requires_grad) {
+      *[](Tensor& self, bool requires_grad) {
         return self.set_requires_grad(requires_grad);
       })
     .define_method(
       "_backward",
-      *[](torch::Tensor& self) {
-        return self.backward();
-      })
-    .define_method(
-      "_backward_gradient",
-      *[](torch::Tensor& self, const torch::Tensor& gradient) {
-        return self.backward(gradient);
+      *[](Tensor& self, Object gradient) {
+        return gradient.is_nil() ? self.backward() : self.backward(from_ruby<torch::Tensor>(gradient));
       })
     .define_method(
       "grad",
-      *[](torch::Tensor& self) {
+      *[](Tensor& self) {
         return self.grad();
       })
     .define_method(
       "_dtype",
-      *[](torch::Tensor& self) {
+      *[](Tensor& self) {
         return (int) at::typeMetaToScalarType(self.dtype());
       })
     .define_method(
       "_type",
-      *[](torch::Tensor& self, int dtype) {
+      *[](Tensor& self, int dtype) {
         return self.toType((torch::ScalarType) dtype);
       })
     .define_method(
       "_layout",
-      *[](torch::Tensor& self) {
+      *[](Tensor& self) {
         std::stringstream s;
         s << self.layout();
         return s.str();
       })
     .define_method(
       "device",
-      *[](torch::Tensor& self) {
+      *[](Tensor& self) {
         std::stringstream s;
         s << self.device();
         return s.str();
       })
     .define_method(
-      "_view",
-      *[](torch::Tensor& self, IntArrayRef size) {
-        return self.view(size);
-      })
-    .define_method(
-      "resize_as!",
-      *[](torch::Tensor& self, torch::Tensor& other) {
-        return self.resize_as_(other);
-      })
-    .define_method(
-      "fill!",
-      *[](torch::Tensor& self, Scalar value) {
-        return self.fill_(value);
-      })
-    .define_method(
-      "_add!",
-      *[](torch::Tensor& self, torch::Tensor& other) {
-        return self.add_(other);
-      })
-    .define_method(
-      "_add_alpha!",
-      *[](torch::Tensor& self, torch::Tensor& other, Scalar alpha) {
-        return self.add_(other, alpha);
-      })
-    .define_method(
-      "_add_scalar!",
-      *[](torch::Tensor& self, Scalar other) {
-        return self.add_(other);
-      })
-    .define_method(
-      "normal!",
-      *[](torch::Tensor& self, double mean, double std) {
-        return self.normal_(mean, std);
-      })
-    .define_method(
-      "sub!",
-      *[](torch::Tensor& self, torch::Tensor& other) {
-        return self.sub_(other);
-      })
-    .define_method(
-      "_mul!",
-      *[](torch::Tensor& self, torch::Tensor& other) {
-        return self.mul_(other);
-      })
-    .define_method(
-      "_mul_scalar!",
-      *[](torch::Tensor& self, Scalar other) {
-        return self.mul_(other);
-      })
-    .define_method(
-      "div!",
-      *[](torch::Tensor& self, torch::Tensor& other) {
-        return self.div_(other);
-      })
-    .define_method(
-      "sqrt!",
-      *[](torch::Tensor& self) {
-        return self.sqrt_();
-      })
-    .define_method(
-      "unsqueeze!",
-      *[](torch::Tensor& self, int64_t dim) {
-        return self.unsqueeze_(dim);
-      })
-    .define_method(
-      "copy!",
-      *[](torch::Tensor& self, torch::Tensor& src) {
-        return self.copy_(src);
-      })
-    .define_method(
-      "clone",
-      *[](torch::Tensor& self) {
-        return self.clone();
-      })
-    .define_method(
-      "log_softmax",
-      *[](torch::Tensor& self, int64_t dim) {
-        return self.log_softmax(dim);
-      })
-    .define_method(
-      "data",
-      *[](torch::Tensor& self) {
-        return self.data();
-      })
-    .define_method(
-      "_data",
-      *[](torch::Tensor& self) {
+      "_flat_data",
+      *[](Tensor& self) {
         Array a;
         auto dtype = self.dtype();
 
@@ -730,23 +262,18 @@ void Init_ext()
         }
         return a;
       })
-    .define_method(
-      "_size",
-      *[](torch::Tensor& self, int i) {
-        return self.size(i);
-      })
     .define_method(
       "_to",
-      *[](torch::Tensor& self, torch::Device device, int dtype, bool non_blocking, bool copy) {
+      *[](Tensor& self, torch::Device device, int dtype, bool non_blocking, bool copy) {
         return self.to(device, (torch::ScalarType) dtype, non_blocking, copy);
       })
     .define_singleton_method(
       "_make_subclass",
-      *[](torch::Tensor& rd, bool requires_grad) {
+      *[](Tensor& rd, bool requires_grad) {
         auto data = torch::autograd::as_variable_ref(rd).detach();
         data.unsafeGetTensorImpl()->set_allow_tensor_metadata_change(true);
         auto var = data.set_requires_grad(requires_grad);
-        return torch::autograd::Variable(std::move(var));
+        return Parameter(std::move(var));
       });
 
   Class rb_cTensorOptions = define_class_under<torch::TensorOptions>(rb_mTorch, "TensorOptions")
@@ -789,36 +316,84 @@ void Init_ext()
         return self.requires_grad(requires_grad);
       });
 
-  Module rb_mNN = define_module_under(rb_mTorch, "NN");
-
   Module rb_mInit = define_module_under(rb_mNN, "Init")
     .define_singleton_method(
-      "kaiming_uniform!",
-      *[](torch::Tensor& input, double a) {
-        return torch::nn::init::kaiming_uniform_(input, a);
+      "_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(
-      "normal!",
-      *[](torch::Tensor& input) {
-        return torch::nn::init::normal_(input);
+      "_ones!",
+      *[](Tensor tensor) {
+        return torch::nn::init::ones_(tensor);
       })
     .define_singleton_method(
-      "uniform!",
-      *[](torch::Tensor& input, double to, double from) {
-        return torch::nn::init::uniform_(input, to, from);
+      "_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);
       });
 
-  Class rb_cParameter = define_class_under<torch::autograd::Variable, torch::Tensor>(rb_mNN, "Parameter")
-    // TODO return grad or nil to remove need for 2nd function
-    .define_method(
-      "_grad",
-      *[](torch::autograd::Variable& self) {
-        return self.grad();
-      })
+  Class rb_cParameter = define_class_under<Parameter, torch::Tensor>(rb_mNN, "Parameter")
     .define_method(
-      "_grad_defined",
-      *[](torch::autograd::Variable& self) {
-        return self.grad().defined();
+      "grad",
+      *[](Parameter& self) {
+        auto grad = self.grad();
+        return grad.defined() ? to_ruby<torch::Tensor>(grad) : Nil;
       });
 
   Class rb_cDevice = define_class_under<torch::Device>(rb_mTorch, "Device")