/aosp_15_r20/external/pytorch/test/ |
H A D | test_nestedtensor.py | 18 from torch.nested._internal.nested_tensor import ( 72 yield torch.nested.nested_tensor 105 nt_contiguous = torch.nested.nested_tensor(ys) 108 nt_noncontiguous = torch.nested.nested_tensor(xs).transpose(-1, -2) 180 return torch.nested.nested_tensor(ts1, device=device, dtype=dtype, layout=layout) 196 return torch.nested.nested_tensor( 224 return torch.nested.nested_tensor( 308 nested_tensor = torch.nested.nested_tensor(data, dtype=torch.int64) 309 nested_tensor_list = nested_tensor.unbind() 330 nested_tensor = torch.nested.nested_tensor(data, dtype=torch.int64) [all …]
|
H A D | test_native_mha.py | 43 x = torch.nested.nested_tensor(xs, device=device, dtype=dtype) 207 q = torch.nested.nested_tensor(qs, device=device, dtype=dtype) 211 k = torch.nested.nested_tensor(torch.unbind(k), device=device, dtype=dtype) 214 k = torch.nested.nested_tensor(torch.unbind(k), device=device, dtype=dtype) 215 v = torch.nested.nested_tensor(torch.unbind(v), device=device, dtype=dtype)
|
H A D | test_transformers.py | 161 return torch.nested.nested_tensor([ 166 return torch.nested.nested_tensor([ 1112 … nested_tensor_return_value = torch.nested.nested_tensor([torch.ones((2, 2), dtype=torch.float)]) 2346 return (torch.nested.nested_tensor(tensors, device=device, dtype=torch.float32), 2347 torch.nested.nested_tensor(tensors, device=device, dtype=torch.float16)) 3234 result = torch.nested.nested_tensor( 3238 … result = torch.nested.nested_tensor([t[0] for _ in range(batch)], dtype=torch.float32) 3241 result = torch.nested.nested_tensor([x.expand(-1, num_heads, t.size(-1)) 3281 …query_expanded = torch.nested.nested_tensor([query.squeeze(0) for _ in range(batch)]).transpose(1,… 3283 value_expanded = torch.nested.nested_tensor( [all …]
|
/aosp_15_r20/external/pytorch/docs/source/ |
H A D | nested.rst | 34 …ion is straightforward and involves passing a list of Tensors to the ``torch.nested.nested_tensor`` 42 >>> nt = torch.nested.nested_tensor([a, b]) 44 nested_tensor([ 51 >>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32, device="cuda", requires_grad=True) 53 nested_tensor([ 66 >>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32) 74 >>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32) 77 RuntimeError: All Tensors given to nested_tensor must have the same dimension. Found dimension 3 fo… 95 >>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32) 108 >>> nt = torch.nested.nested_tensor([a, a], dtype=torch.float32) [all …]
|
/aosp_15_r20/external/pytorch/torch/nested/ |
H A D | __init__.py | 121 from torch.nested._internal.nested_tensor import nested_view_from_values_offsets 127 from torch.nested._internal.nested_tensor import jagged_from_list 194 def nested_tensor(tensor_list, *, dtype=None, layout=None, device=None, requires_grad=False, pin_me… function 226 return _nested.nested_tensor( 236 from torch.nested._internal.nested_tensor import jagged_from_list 301 from torch.nested._internal.nested_tensor import jagged_from_tensor_and_lengths 410 from torch.nested._internal.nested_tensor import nested_view_from_values_offsets_lengths 458 from torch.nested._internal.nested_tensor import (
|
/aosp_15_r20/external/pytorch/aten/src/ATen/native/nested/ |
H A D | NestedTensorTransformerFunctions.cpp | 16 const Tensor& nested_tensor, in check_nested_tensor_matrix_constraints() argument 19 auto* nt_input = get_nested_tensor_impl(nested_tensor); in check_nested_tensor_matrix_constraints() 31 nested_tensor.dim() == 3 && dense_matrix.dim() == 2, in check_nested_tensor_matrix_constraints() 35 nested_tensor.dim(), in check_nested_tensor_matrix_constraints()
|
/aosp_15_r20/external/pytorch/benchmarks/nested/ |
H A D | nested_bmm_bench.py | 26 nt_a = torch.nested.nested_tensor( 31 nt_b = torch.nested.nested_tensor(
|
/aosp_15_r20/external/pytorch/torch/csrc/api/include/torch/ |
H A D | nested.h | 23 inline at::Tensor nested_tensor( 38 inline at::Tensor nested_tensor(
|
/aosp_15_r20/external/pytorch/benchmarks/transformer/ |
H A D | better_transformer_vs_mha_functional.py | 84 q = torch.nested.nested_tensor(q, device=device, dtype=dtype) 100 torch.nested.nested_tensor(
|
H A D | sdp.py | 182 torch.nested.nested_tensor(
|
/aosp_15_r20/external/pytorch/torch/autograd/ |
H A D | __init__.py | 66 from torch.nested._internal.nested_tensor import NestedTensor 112 from torch.nested._internal.nested_tensor import NestedTensor
|
/aosp_15_r20/external/pytorch/test/cpp/api/ |
H A D | nested.cpp | 13 auto nt = torch::nested::nested_tensor({a, b}); in TEST()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/lite/ |
H A D | flatbuffer_import.cc | 222 for (const auto& nested_tensor : tensor.variant_tensors) { in GetTensorType() local 224 ConvertElementType(nested_tensor->type, builder); in GetTensorType() 225 if (nested_tensor->has_rank) { in GetTensorType() 226 llvm::SmallVector<int64_t> shape(nested_tensor->shape.begin(), in GetTensorType() 227 nested_tensor->shape.end()); in GetTensorType()
|
/aosp_15_r20/external/pytorch/torch/utils/ |
H A D | _device.py | 33 torch.nested.nested_tensor,
|
/aosp_15_r20/external/pytorch/torch/nested/_internal/ |
H A D | ops.py | 12 from .nested_tensor import NestedTensor 445 from .nested_tensor import jagged_from_list 506 from .nested_tensor import _tensor_symint_registry 1667 from torch.nested._internal.nested_tensor import _nt_view_dummy
|
H A D | sdpa.py | 18 from .nested_tensor import buffer_from_jagged, NestedTensor, ViewNestedFromBuffer
|
/aosp_15_r20/external/pytorch/test/profiler/ |
H A D | test_profiler.py | 922 nested_tensor = torch.nested.nested_tensor( 930 nested_tensor = nested_tensor + nested_tensor 1421 inp = torch.nested.nested_tensor([a, b])
|
H A D | test_profiler_tree.py | 927 import torch.nested._internal.nested_tensor
|
/aosp_15_r20/external/pytorch/torch/_subclasses/ |
H A D | fake_tensor.py | 1181 import torch.nested._internal.nested_tensor 1184 torch.nested._internal.nested_tensor._tensor_id_counter 1245 import torch.nested._internal.nested_tensor 2150 import torch.nested._internal.nested_tensor
|
/aosp_15_r20/external/pytorch/torch/testing/_internal/opinfo/definitions/ |
H A D | nested.py | 38 return torch.nested.nested_tensor(
|
/aosp_15_r20/external/pytorch/test/nn/ |
H A D | test_multihead_attention.py | 707 nt = torch.nested.nested_tensor([torch.randn(4, 4)]) 737 nt = torch.nested.nested_tensor([torch.randn(4, 4, requires_grad=False)])
|
/aosp_15_r20/external/pytorch/test/dynamo/ |
H A D | test_subclasses.py | 20 from torch.nested._internal.nested_tensor import ( 1868 from torch.nested._internal.nested_tensor import _tensor_symint_registry 1872 count_before = torch.nested._internal.nested_tensor._tensor_id_counter 1892 count_after = torch.nested._internal.nested_tensor._tensor_id_counter
|
/aosp_15_r20/external/pytorch/torch/multiprocessing/ |
H A D | reductions.py | 328 from torch.nested._internal.nested_tensor import NestedTensor
|
/aosp_15_r20/external/pytorch/torch/_dynamo/variables/ |
H A D | torch.py | 703 @register(torch.nested.nested_tensor)
|
H A D | builder.py | 1566 and not isinstance(value, torch.nested._internal.nested_tensor.NestedTensor) 2365 e, torch.nested._internal.nested_tensor.NestedTensor
|