Home
last modified time | relevance | path

Searched refs:nested_tensor (Results 1 – 25 of 33) sorted by relevance

12

/aosp_15_r20/external/pytorch/test/
H A Dtest_nestedtensor.py18 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 Dtest_native_mha.py43 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 Dtest_transformers.py161 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 Dnested.rst34 …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__.py121 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 DNestedTensorTransformerFunctions.cpp16 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 Dnested_bmm_bench.py26 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 Dnested.h23 inline at::Tensor nested_tensor(
38 inline at::Tensor nested_tensor(
/aosp_15_r20/external/pytorch/benchmarks/transformer/
H A Dbetter_transformer_vs_mha_functional.py84 q = torch.nested.nested_tensor(q, device=device, dtype=dtype)
100 torch.nested.nested_tensor(
H A Dsdp.py182 torch.nested.nested_tensor(
/aosp_15_r20/external/pytorch/torch/autograd/
H A D__init__.py66 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 Dnested.cpp13 auto nt = torch::nested::nested_tensor({a, b}); in TEST()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/lite/
H A Dflatbuffer_import.cc222 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.py33 torch.nested.nested_tensor,
/aosp_15_r20/external/pytorch/torch/nested/_internal/
H A Dops.py12 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 Dsdpa.py18 from .nested_tensor import buffer_from_jagged, NestedTensor, ViewNestedFromBuffer
/aosp_15_r20/external/pytorch/test/profiler/
H A Dtest_profiler.py922 nested_tensor = torch.nested.nested_tensor(
930 nested_tensor = nested_tensor + nested_tensor
1421 inp = torch.nested.nested_tensor([a, b])
H A Dtest_profiler_tree.py927 import torch.nested._internal.nested_tensor
/aosp_15_r20/external/pytorch/torch/_subclasses/
H A Dfake_tensor.py1181 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 Dnested.py38 return torch.nested.nested_tensor(
/aosp_15_r20/external/pytorch/test/nn/
H A Dtest_multihead_attention.py707 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 Dtest_subclasses.py20 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 Dreductions.py328 from torch.nested._internal.nested_tensor import NestedTensor
/aosp_15_r20/external/pytorch/torch/_dynamo/variables/
H A Dtorch.py703 @register(torch.nested.nested_tensor)
H A Dbuilder.py1566 and not isinstance(value, torch.nested._internal.nested_tensor.NestedTensor)
2365 e, torch.nested._internal.nested_tensor.NestedTensor

12