From 14ead627f2858ee2dab52bbb53208054319a81ef Mon Sep 17 00:00:00 2001
From: min-jean-cho <min.jean.cho@intel.com>
Date: Thu, 9 Jan 2025 00:45:34 -0800
Subject: [PATCH] Add device-agnostic NestedTensor Ops  (#1258)

Device-agnostic ops. CUDA, CPU dispatched to the same ops as well.

- `_nested_compute_contiguous_strides_offsets`
- `_nested_tensor_from_mask`
- `_nested_tensor_from_mask_left_aligned`
- `_nested_view_from_buffer`
---
 yaml/native/native_functions.yaml | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/yaml/native/native_functions.yaml b/yaml/native/native_functions.yaml
index db6821026..a1dc859ed 100644
--- a/yaml/native/native_functions.yaml
+++ b/yaml/native/native_functions.yaml
@@ -8612,3 +8612,32 @@
     XPU: _weight_int4pack_mm_xpu
   # autogen: _weight_int4pack_mm.out
   # tags: core
+
+  - func: _nested_compute_contiguous_strides_offsets(Tensor nested_size) -> (Tensor, Tensor)
+  variants: function
+  device_check: NoCheck
+  dispatch:
+    XPU: _nested_compute_contiguous_strides_offsets
+
+- func: _nested_tensor_from_mask(Tensor t, Tensor mask, bool mask_check=True) -> Tensor
+  dispatch:
+    XPU: NestedTensor_nested_tensor_from_mask
+  autogen: _nested_tensor_from_mask.out
+
+- func: _nested_tensor_from_mask_left_aligned(Tensor t, Tensor mask) -> bool
+  dispatch:
+    XPU: NestedTensor_nested_tensor_from_mask_left_aligned
+
+- func: _nested_view_from_buffer(Tensor(a) self, Tensor nested_size, Tensor nested_strides, Tensor offsets) -> Tensor(a)
+  variants: function
+  device_check: NoCheck
+  dispatch:
+    XPU: _nested_view_from_buffer
+
+- func: _nested_view_from_buffer_copy(Tensor self, Tensor nested_size, Tensor nested_strides, Tensor offsets) -> Tensor
+  variants: function
+  device_check: NoCheck
+  tags: view_copy
+  dispatch:
+    CompositeExplicitAutogradNonFunctional: _nested_view_from_buffer_copy
+  autogen: _nested_view_from_buffer_copy.out
\ No newline at end of file