Repository: pytorch/functorch (language: python)

Issue:
Can't call `torch.tensor` within grad
```
from functorch import grad, vmap
import torch

def f(x):
    t = torch.tensor(0)
    return t + x
inps = (torch.randn([]),)
print(grad(f)(*inps))
```

```
Traceback (most recent call last):
  File "t.py", line 8, in <module>
    print(grad(f)(*inps))
  File "/home/chilli/fb/functorch/functorch/_src/eager_transforms.py", line 152, in wrapper
    results = grad_and_value(f, argnums, has_aux=has_aux)(*args)
  File "/home/chilli/fb/functorch/functorch/_src/eager_transforms.py", line 110, in wrapper
    output = f(*args)
  File "t.py", line 5, in f
    t = torch.tensor(0)
RuntimeError: Cannot access data pointer of Tensor that doesn't have storage
```


Discussion / hints:
https://github.com/zou3519/functorch/issues/14 is a pretty fun test case. One of my previous proposed solutions, "whenever you call empty, don't return a TensorWrapper; return a TensorWrapper for all other operations" doesn't actually work. The problem is that if the user calls empty and then does a .copy_(), we really wanted the tensor returned from empty to have a TensorWrapper because the .copy_() is going to assign history to it.
I think we're going to need an extension point in core for this. For the other solution we discussed (just attaching the storage to TensorWrapper) I'm worried about aliasing because this would effectively make TensorWrapper a view on another Tensor.
One of the problems with `torch.tensor` is we're reading in the data from a legitimately arbitrarily Python object, so there isn't really any way to get things to work without mutating a tensor first. What it seems like you want to happen in this case is for a normal, non mode overridden tensor to get constructed, you mutate into it, and THEN you alias it into the wrapper.
>  What it seems like you want to happen in this case is for a normal, non mode overridden tensor to get constructed, you mutate into it, and THEN you alias it into the wrapper.

Here are two possible options:

## Option 1: Easy fix

To achieve this it looks like we can just put ` c10::impl::ExcludeDispatchKeyGuard dynamic_layer_back_guard(c10::DispatchKey::FuncTorchDynamicLayerBackMode);`  [here](https://github.com/pytorch/pytorch/blob/7f6b2bc2d01cb94cbd24b563479a365d13b100e9/torch/csrc/utils/tensor_new.cpp#L260-L261)

The FuncTorchDynamicLayerBackMode key's fallback is responsible for re-wrapping tensors and we want to make it so that the normal tensor gets constructed and we mutate into it. Aliasing in the wrapper happens in the call to `.to(...)` ([here](https://github.com/pytorch/pytorch/blob/7f6b2bc2d01cb94cbd24b563479a365d13b100e9/torch/csrc/utils/tensor_new.cpp#L278)) because `.to` will dispatch to  FuncTorchDynamicLayerBackMode

## Option 2: Add an extension point for `internal_new_from_data`

We introduce some operator that does `internal_new_from_data` (https://github.com/pytorch/pytorch/blob/7f6b2bc2d01cb94cbd24b563479a365d13b100e9/torch/csrc/utils/tensor_new.cpp#L201).

A glaring problem is that native_functions codegen doesn't support PyObject* data. But this solution would be more general.
We don't have a particular use case for this generic solution yet, so I'm inclined to go with Option 1 for now

Fix patch:
diff --git a/functorch/csrc/TensorWrapper.cpp b/functorch/csrc/TensorWrapper.cpp
index 7a2745aab..dcc95928d 100644
--- a/functorch/csrc/TensorWrapper.cpp
+++ b/functorch/csrc/TensorWrapper.cpp
@@ -143,11 +143,28 @@ TensorWrapper::TensorWrapper(
   , is_alive_(std::move(is_alive))
 {
   TORCH_INTERNAL_ASSERT(value_.defined());
-  set_storage_access_should_throw();
 
   // TODO: need to reset sizes/strides on mutation
   TORCH_INTERNAL_ASSERT(use_value_sizes_strides);
   refreshMetadata();
+
+  // Note [TensorWrapper sometimes has storage]
+  // TensorWrapper has storage if it isn't wrapping a BatchedTensor.
+  // If TensorWrapper ultimately wraps a BatchedTensor (directly or
+  // indirectly), then it does not have storage.
+  //
+  // TensorImpl's storage_access_should_throw_ field is private
+  // (TODO: we probably want to make it protected or add an accessor)
+  // so we work around it by trying to figure out if the storage access
+  // will throw or not. A possible alternative is to try-catch .storage().
+  auto* wrapped = maybeGetTensorWrapper(value_);
+  auto* batched = maybeGetBatchedImpl(value_);
+  if (batched || wrapped && wrapped->storage_access_should_throw()) {
+    hacky_storage_access_should_throw_ = true;
+    set_storage_access_should_throw();
+  } else {
+    storage_ = value_.storage();
+  }
 }
 
 // The following are some internal inherited methods that we do not support.
diff --git a/functorch/csrc/TensorWrapper.h b/functorch/csrc/TensorWrapper.h
index 767ee0e82..852093630 100644
--- a/functorch/csrc/TensorWrapper.h
+++ b/functorch/csrc/TensorWrapper.h
@@ -46,6 +46,11 @@ struct TORCH_API TensorWrapper : public c10::TensorImpl {
       bool allow_tensor_metadata_change) const override;
   void shallow_copy_from(const c10::intrusive_ptr<TensorImpl>& impl) override;
 
+  // See Note [TensorWrapper sometimes has storage]
+  bool storage_access_should_throw() {
+    return hacky_storage_access_should_throw_;
+  }
+
  private:
   const char* tensorimpl_type_name() const override;
   Tensor value_;
@@ -56,6 +61,9 @@ struct TORCH_API TensorWrapper : public c10::TensorImpl {
   // 1) May still have autograd metadata on them
   // 2) Forward dispatches to the underlying value()
   std::shared_ptr<bool> is_alive_;
+
+  // Note [TensorWrapper sometimes has storage]
+  bool hacky_storage_access_should_throw_ = false;
 };
 
 TORCH_API Tensor makeTensorWrapper(const Tensor& tensor, int64_t level);


Test patch:
diff --git a/test/test_eager_transforms.py b/test/test_eager_transforms.py
index 6fff09e1c..1fe76fc39 100644
--- a/test/test_eager_transforms.py
+++ b/test/test_eager_transforms.py
@@ -463,7 +463,6 @@ def f2(value):
         result = grad(f2)(x)
         self.assertEqual(result, (x <= 0).type_as(x))
 
-    @unittest.expectedFailure
     def test_tensor_ctor_inside_grad(self, device):
         def foo(x):
             return x * torch.tensor(2., device=device)
diff --git a/test/test_ops.py b/test/test_ops.py
index 26b7c161c..d16a4b344 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -136,12 +136,10 @@ def is_inplace(op, variant):
 
 
 vjp_fail = {
-    '__getitem__',
-    '__rpow__',
     'linalg.cholesky',
     'linalg.inv',
-    'linalg.matrix_norm',
     'linalg.matrix_power',
+    'linalg.matrix_norm',
     'linalg.norm',
     'nanquantile',
     'quantile',
@@ -216,7 +214,8 @@ def test_vjp(self, device, dtype, op):
             result = fn(*primals)
             cotangents = tree_map(lambda x: torch.randn_like(x), result)
 
-            _, vjp_fn = vjp(fn, *primals)
+            out, vjp_fn = vjp(fn, *primals)
+            self.assertEqual(out, result)
             result_vjps = vjp_fn(cotangents)
 
             _, vjp_fn = ref_vjp(fn, *primals)