Home
last modified time | relevance | path

Searched refs:NestedTensor (Results 1 – 19 of 19) sorted by relevance

/aosp_15_r20/external/pytorch/torch/nested/_internal/
H A Dops.py12 from .nested_tensor import NestedTensor
84 "t": lambda x: isinstance(x, torch.Tensor) and not isinstance(x, NestedTensor),
85 "jt": lambda x: isinstance(x, NestedTensor)
89 x, NestedTensor
132 func, a: NestedTensor, a_name: str, b: NestedTensor, b_name: str argument
244 njt = next(arg for arg in args if isinstance(arg, NestedTensor))
245 return NestedTensor(
253 assert isinstance(a, NestedTensor) or isinstance(b, NestedTensor)
259 if isinstance(a, NestedTensor) and isinstance(b, NestedTensor):
263 return NestedTensor(
[all …]
H A Dnested_tensor.py48 class NestedTensor(torch.Tensor): class
81 ks = DispatchKeySet(DispatchKey.NestedTensor)
87 assert not isinstance(values, NestedTensor)
231 func = NestedTensor
275 return NestedTensor(
325 def forward(ctx, x: NestedTensor): # type: ignore[override] argument
334 return NestedTensor(
360 return NestedTensor(
367 def backward(ctx, gO: NestedTensor): # type: ignore[override] argument
381 ) -> Tuple[NestedTensor, torch.Tensor]:
[all …]
H A Dsdpa.py18 from .nested_tensor import buffer_from_jagged, NestedTensor, ViewNestedFromBuffer
33 not isinstance(query, NestedTensor)
34 or not isinstance(key, NestedTensor)
35 or not isinstance(value, NestedTensor)
116 assert isinstance(param, NestedTensor), "param should be a jagged NT"
312 if not isinstance(qkv, NestedTensor):
342 assert isinstance(tensor, NestedTensor)
629 isinstance(query, NestedTensor)
630 and isinstance(key, NestedTensor)
631 and isinstance(value, NestedTensor)
[all …]
/aosp_15_r20/external/pytorch/aten/src/ATen/native/nested/
H A DREADME.md4 ## NestedTensor Data Structure
6 …sions. The tensor_components are flattened and combined into a single NestedTensor, which includes…
15 …NestedTensor in C++ you will likely not be using the NestedTensorImpl constructor directly but usi…
19 The NestedTensor code is split into two parts: the C++ code and the Python code. The C++ code is lo…
21 - `NestedTensorImpl.h | NestedTensorImpl.cpp`: The NestedTensor data structure and its methods.
29 - `cuda/`: CUDA implementations of the NestedTensor functions
32 …Tensors: functions that can be efficiently implemented by viewing the NestedTensor as dense tensor…
44 1. Get the NestedTensorImpl from the input NestedTensor.
45 2. Get the sizes of the NestedTensor.
46 …et_buffer() which returns a flat, dense tensor whose storage shares that of the input NestedTensor.
[all …]
/aosp_15_r20/external/pytorch/c10/core/
H A DDispatchKeySet.cpp42 && t != DispatchKey::NestedTensor && backend_dispatch_keyset.has(t); in isBackendDispatchKey()
55 DispatchKeySet{DispatchKey::NestedTensor} |
61 {DispatchKey::AutogradNestedTensor, DispatchKey::NestedTensor}) |
100 return k != DispatchKey::NestedTensor && backend_dispatch_keyset.has(k); in runtimeDispatchKeySetHas()
103 return k != DispatchKey::NestedTensor && in runtimeDispatchKeySetHas()
141 return DispatchKeySet(DispatchKey::NestedTensor) | in getBackendKeySetFromAutograd()
H A DDispatchKey.h61 _(NestedTensor, NestedTensor) \
221 NestedTensor, enumerator
551 k == DispatchKey::NestedTensor) { in isPerBackendFunctionalityKey()
673 return DispatchKey::NestedTensor; in toFunctionalityKey()
712 if (functionality_k == DispatchKey::NestedTensor) { in toRuntimePerBackendFunctionalityKey()
H A DDispatchKey.cpp98 case DispatchKey::NestedTensor: in toString()
328 {"NestedTensor", c10::DispatchKey::NestedTensor}, in parseDispatchKey()
H A DTensorImpl.h1218 return key_set_.has(DispatchKey::NestedTensor); in is_nested()
/aosp_15_r20/external/pytorch/docs/source/
H A Dnested.rst13 NestedTensor allows the user to pack a list of Tensors into a single, efficient datastructure.
20 While the conventional approach is to pad variable length sequences, NestedTensor
62 In order to form a valid NestedTensor all the passed Tensors need to match in dimension, but none o…
80 NestedTensor allocates new memory to store them and does not keep a reference.
84 …is important to maintain an even level of nesting across entries so that the resulting NestedTensor
91 Even though a NestedTensor does not support ``.size()`` (or ``.shape``), it supports ``.size(i)`` i…
105 If all dimensions are regular, the NestedTensor is intended to be semantically indistinguishable fr…
162 …lice of the underlying memory, which represents the first entry or constituent of the NestedTensor.
183 NestedTensor and any constraints they have.
222 …ed_shape`` argument is restricted to not extend into the irregular dimensions of the NestedTensor."
/aosp_15_r20/external/pytorch/torch/autograd/
H A D__init__.py66 from torch.nested._internal.nested_tensor import NestedTensor
75 if output.is_nested and not isinstance(output, NestedTensor):
112 from torch.nested._internal.nested_tensor import NestedTensor
117 out_is_cpp_nested = out_is_nested and not isinstance(out, NestedTensor)
/aosp_15_r20/external/pytorch/torch/multiprocessing/
H A Dreductions.py328 from torch.nested._internal.nested_tensor import NestedTensor
330 if tensor.is_nested and not isinstance(tensor, NestedTensor):
/aosp_15_r20/external/pytorch/aten/src/ATen/
H A DNestedTensorImpl.h196 c10::DispatchKey::NestedTensor, in generate_buffer_key_set()
H A DNestedTensorImpl.cpp50 nested_key_set | c10::DispatchKeySet{c10::DispatchKey::NestedTensor}; in generate_nested_key_set_from_buffer()
/aosp_15_r20/external/pytorch/torch/nn/modules/
H A Dactivation.py1270 is not supported with NestedTensor input"
/aosp_15_r20/external/pytorch/test/
H A Dtest_nestedtensor.py22 NestedTensor,
5713 nt = torch.nested._internal.nested_tensor.NestedTensor(
5732 nt = torch.nested._internal.nested_tensor.NestedTensor(
5747 nt = torch.nested._internal.nested_tensor.NestedTensor(
5766 nt = torch.nested._internal.nested_tensor.NestedTensor(
5788 nt = torch.nested._internal.nested_tensor.NestedTensor(
6540 self.assertTrue(isinstance(output, NestedTensor))
H A Dtest_python_dispatch.py1313 DispatchKeySet(DispatchKey.NestedTensor),
1322 self.assertTrue(torch._C._dispatch_keys(x).has(DispatchKey.NestedTensor))
/aosp_15_r20/external/pytorch/docs/source/community/
H A Dpersons_of_interest.rst167 NestedTensor (torch.nested)
/aosp_15_r20/external/pytorch/torch/_dynamo/variables/
H A Dbuilder.py1566 and not isinstance(value, torch.nested._internal.nested_tensor.NestedTensor)
2365 e, torch.nested._internal.nested_tensor.NestedTensor
/aosp_15_r20/external/pytorch/torchgen/
H A Dmodel.py95 NestedTensor = auto() variable in DispatchKey