pytorch · jbschlosser · Oct 29, 2024 · Oct 29, 2024 · Oct 30, 2024 · Oct 30, 2024
diff --git a/aten/src/ATen/FunctionalInverses.cpp b/aten/src/ATen/FunctionalInverses.cpp
@@ -465,6 +465,14 @@ Tensor FunctionalInverses::narrow_inverse(const at::Tensor & base, const at::Ten
     }
 }
 
+Tensor FunctionalInverses::view_as_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, const at::Tensor & other) {
+    if (inverse_return_mode != InverseReturnMode::NeverView) {
+      return mutated_view.view_as(base);
+    } else {
+      return mutated_view.view_as(base).clone();
+    }
+}
+
 Tensor FunctionalInverses::slice_inverse_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, const at::Tensor & src, int64_t dim, std::optional<c10::SymInt> start, std::optional<c10::SymInt> end, c10::SymInt step) {
     // slice_inverse() inverse is just slice()
     if (inverse_return_mode == InverseReturnMode::NeverView) {

diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -6475,6 +6475,9 @@
   variants: method
   device_check: NoCheck
   device_guard: False
+  dispatch:
+    CompositeImplicitAutograd: view_as
+    NestedTensorCPU, NestedTensorCUDA: view_as_nested
 
 - func: where.self(Tensor condition, Tensor self, Tensor other) -> Tensor
   device_check: NoCheck   # TensorIterator

diff --git a/aten/src/ATen/native/nested/NestedTensorMath.cpp b/aten/src/ATen/native/nested/NestedTensorMath.cpp
@@ -1090,4 +1090,9 @@ Tensor cat_nested(const ITensorListRef& tensors, int64_t dim) {
   return cat_nested_impl(materialized, at::legacy_cat_wrap_dim(dim, materialized));
 }
 
+Tensor view_as_nested(const Tensor& self, const Tensor& other) {
+  TORCH_INTERNAL_ASSERT(false, "view_as(): only implemented for jagged layout nested tensors");
+  return Tensor();
+}
+
 } // namespace at::native
diff --git a/aten/src/ATen/templates/FunctionalInverses.h b/aten/src/ATen/templates/FunctionalInverses.h
@@ -27,6 +27,7 @@ struct FunctionalInverses {
 // https://github.com/pytorch/pytorch/blob/main/torchgen/model.py#L2583-L2585
 static at::Tensor chunk_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, int64_t mutated_view_idx, int chunks, int dim);
 static at::Tensor narrow_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, int dim, c10::SymInt start, c10::SymInt length);
+static at::Tensor view_as_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, const at::Tensor & other);
 
 };
 }

diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
@@ -1932,6 +1932,16 @@
   self: at::view_as_real(grad.contiguous().resolve_conj()) # [gx, gy]
   result: at::view_as_complex(self_t)
 
+- name: view_as(Tensor(a) self, Tensor other) -> Tensor(a)
+  dispatch:
+    Default:
+      # the default case will use the CompositeImplicitAutograd impl
+      self: not_implemented("view_as")
+      other: non_differentiable
+    AutogradNestedTensor:
+      self: grad.view_as(self)
+      other: non_differentiable
+
 - name: where.self(Tensor condition, Tensor self, Tensor other) -> Tensor
   condition: non_differentiable
   self: where(condition, grad, 0)

diff --git a/torch/nested/_internal/ops.py b/torch/nested/_internal/ops.py
@@ -1458,6 +1458,25 @@ def get_inner_size(inner_idx):
         return NestedTensor(func(inp._values, inner_size), **extract_kwargs(inp))
 
 
+@register_jagged_func([torch.ops.aten.view_as.default], "self: jt, other: jt")
+def view_as_default(func, *args, **kwargs):
+    _, new_kwargs = normalize_function(  # type: ignore[misc]
+        func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True
+    )
+
+    inp = new_kwargs.pop("input")
+    other = new_kwargs.pop("other")
+
+    error_message = f"view_as(): Cannot view NJT of shape {inp.shape} as shape {other.shape}"
+
+    # verify input is viewable as other's shape
+    if inp._ragged_idx != other._ragged_idx:
+        raise RuntimeError(error_message)
+    torch._assert_async(torch.all(inp._offsets == other._offsets), error_message)
+
+    return NestedTensor(func(inp._values, other._values), **extract_kwargs(other))
+
+
 @register_jagged_func(
     torch.ops.aten.native_layer_norm.default,
     "input: jt_all, normalized_shape: any, weight: any?, bias: any?, eps: any",