Tensor Indexing API for set tensor (#174)

Summary:
Pull Request resolved: #174

Implement Tensor Indexing API for inplace put `index_put_`, which will allow the following notions:

```javascript
const tensor = torch.zeros([2]);
tensor[0] = 1;
tensor[1] = torch.tensor([2]);
```

re #172

Reviewed By: justinhaaheim

Differential Revision: D41746329

fbshipit-source-id: b49a207d4c70f551eb1288bea68e09671c389ae4

Author: raedle

react-native-pytorch-core/cxx/src/torchlive/torch/TensorHostObject.cpp

@@ -737,6 +737,39 @@ jsi::Value TensorHostObject::get(
   return BaseHostObject::get(runtime, propNameId);
 }
 
+void TensorHostObject::set(
+    jsi::Runtime& runtime,
+    const jsi::PropNameID& propNameId,
+    const jsi::Value& value) {
+  auto name = propNameId.utf8(runtime);
+
+  int idx = -1;
+  try {
+    idx = std::stoi(name.c_str());
+  } catch (...) {
+    // Cannot parse name value to int. This can happen when the name in bracket
+    // or dot notion is not an int (e.g., tensor['foo']).
+    // Let's ignore this exception here and have the PyTorch C++ API throw an
+    // error for index out of bounds.
+    // Note: The Tensor Indexing API allows for a much broader range of indices
+    // but for now, the PlayTorch API only supports single value index values.
+  }
+  if (value.isObject()) {
+    // Get TensorHostObject with wrapped tensor, otherwise it will be nullptr
+    auto tensorHostObject =
+        value.asObject(runtime).asHostObject<TensorHostObject>(runtime);
+    if (tensorHostObject != nullptr) {
+      this->tensor.index_put_({idx}, tensorHostObject->tensor);
+    }
+  } else if (value.isNumber()) {
+    this->tensor.index_put_({idx}, value.asNumber());
+  } else {
+    throw jsi::JSError(
+        runtime,
+        "Invalid value! The value has to be of type tensor or a number");
+  }
+}
+
 jsi::Function TensorHostObject::createToString(jsi::Runtime& runtime) {
   auto toStringFunc = [this](
                           jsi::Runtime& runtime,

react-native-pytorch-core/cxx/src/torchlive/torch/TensorHostObject.h

@@ -38,6 +38,10 @@ class JSI_EXPORT TensorHostObject : public common::BaseHostObject {
   facebook::jsi::Value get(
       facebook::jsi::Runtime&,
       const facebook::jsi::PropNameID& name) override;
+  void set(
+      facebook::jsi::Runtime&,
+      const facebook::jsi::PropNameID& name,
+      const facebook::jsi::Value& value) override;
   std::vector<facebook::jsi::PropNameID> getPropertyNames(
       facebook::jsi::Runtime& rt) override;
 

react-native-pytorch-core/cxx/test/TensorTests.cpp

@@ -360,14 +360,14 @@ TEST_F(TorchliveTensorRuntimeTest, TensorIndexing) {
       )";
   EXPECT_TRUE(eval(tensorAccessWithIndex).getBool());
 
-  std::string nestedTensorAcessWithIndex =
+  std::string nestedTensorAccessWithIndex =
       R"(
         const tensor = torch.tensor([[128], [255]]);
-        const tensor1 = tensor[0];
-        const tensor2 = tensor[1];
-        tensor1[0].item() == 128 && tensor2[0].item() == 255;
+        const tensor0 = tensor[0];
+        const tensor1 = tensor[1];
+        tensor0[0].item() == 128 && tensor1[0].item() == 255;
       )";
-  EXPECT_TRUE(eval(nestedTensorAcessWithIndex).getBool());
+  EXPECT_TRUE(eval(nestedTensorAccessWithIndex).getBool());
 
   EXPECT_TRUE(eval("torch.tensor([[128], [255]])['foo']").isUndefined());
 
@@ -376,6 +376,48 @@ TEST_F(TorchliveTensorRuntimeTest, TensorIndexing) {
   EXPECT_TRUE(eval("torch.tensor([[128], [255]])[2]").isUndefined());
 }
 
+TEST_F(TorchliveTensorRuntimeTest, TensorIndexingPut) {
+  std::string tensorPutWithIndex =
+      R"(
+        const tensor = torch.zeros([3]);
+        tensor[0] = torch.tensor([1]);
+        tensor[1] = torch.tensor([2]);
+        tensor[2] = torch.tensor([3]);
+        tensor[0].item() === 1 && tensor[1].item() === 2 && tensor[2].item() === 3;
+      )";
+  EXPECT_TRUE(eval(tensorPutWithIndex).getBool());
+
+  std::string tensorPutWithIndexAndNumberValue =
+      R"(
+        const tensor = torch.zeros([3]);
+        tensor[0] = 1;
+        tensor[1] = 2;
+        tensor[2] = 3;
+        tensor[0].item() === 1 && tensor[1].item() === 2 && tensor[2].item() === 3;
+      )";
+  EXPECT_TRUE(eval(tensorPutWithIndexAndNumberValue).getBool());
+
+  std::string nestedTensorPutWithIndex =
+      R"(
+        const tensor = torch.tensor([[128], [0]]);
+        tensor[1] = torch.tensor([[255]]);
+        const tensor0 = tensor[0];
+        const tensor1 = tensor[1];
+        tensor0[0].item() == 128 && tensor1[0].item() == 255;
+      )";
+  EXPECT_TRUE(eval(nestedTensorPutWithIndex).getBool());
+
+  EXPECT_THROW(
+      eval("torch.tensor([[128], [255]])['foo'] = 'bar'"),
+      facebook::jsi::JSError);
+
+  EXPECT_THROW(
+      eval("torch.tensor([[128], [255]])[-1] = 'bar'"), facebook::jsi::JSError);
+
+  EXPECT_THROW(
+      eval("torch.tensor([[128], [255]])[2] = 'bar'"), facebook::jsi::JSError);
+}
+
 TEST_F(TorchliveTensorRuntimeTest, TensorDivTest) {
   std::string tensorDivWithNumber =
       R"(