[optim] Add layout=jagged to irregular nested tensors

Co-authored-by: tianyi-huawei<getianyi1@huawei.com>



# message auto-generated for no-merge-commit merge:
!16 merge feat/jagged into main

[optim] Add layout=jagged to irregular nested tensors

Created-by: tianyi-huawei
Commit-by: tianyi-huawei
Merged-by: ascend-robot
Description: add layout=jagged to nested tensors because torch warning suggests this usage

before:
![image.png](https://raw.gitcode.com/user-images/assets/8886051/b1de5079-8366-4f24-a9fb-e1df0320564f/image.png 'image.png')

after:
![image.png](https://raw.gitcode.com/user-images/assets/8886051/9d10c041-245e-4e91-9386-83a66efe448f/image.png 'image.png')

See merge request: Ascend/TransferQueue!16

Author: ascend-robot

transfer_queue/storage/managers/base.py

@@ -396,7 +396,7 @@ def _merge_tensors_to_tensordict(metadata: BatchMeta, values: list[Tensor]) -> T
                 except RuntimeError:
                     try:
                         # Fallback to nested tensor if shapes are irregular
-                        merged_data[field] = torch.nested.as_nested_tensor(data_list)
+                        merged_data[field] = torch.nested.as_nested_tensor(data_list, layout=torch.jagged)
                     except Exception:
                         merged_data[field] = NonTensorStack(*data_list)
             else:

transfer_queue/storage/managers/simple_backend_manager.py

@@ -313,7 +313,7 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict:
                 and all(isinstance(item, torch.Tensor) for item in v)
                 and all(item.shape == v[0].shape for item in v)
                 else (
-                    torch.nested.as_nested_tensor(v)
+                    torch.nested.as_nested_tensor(v, layout=torch.jagged)
                     if v and all(isinstance(item, torch.Tensor) for item in v)
                     else NonTensorStack(*v)
                 )

transfer_queue/utils/serial_utils.py

@@ -264,7 +264,9 @@ def deserialization(data: list[bytestr] | bytestr) -> Any:
                 tensors[i] = single_tensors[current_idx]
                 current_idx += 1
             else:
-                tensors[i] = torch.nested.as_nested_tensor(single_tensors[current_idx : current_idx + tensor_num])
+                tensors[i] = torch.nested.as_nested_tensor(
+                    single_tensors[current_idx : current_idx + tensor_num], layout=torch.strided
+                )
                 current_idx += tensor_num
 
         return _internal_rpc_pickler.deserialize(pickled_bytes, tensors)