Searched refs:_nested_tensor_size (Results 1 – 12 of 12) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/csrc/autograd/ |
| H A D | input_metadata.cpp | 12 auto nested_size = input._nested_tensor_size(); in compute_variant_shape()
122 return grad._nested_tensor_size().is_same_size(shape_as_tensor()); in is_same_shape()
146 ss << grad._nested_tensor_size(); in incompatible_shape_error_message()
|
| H A D | saved_variable.cpp | 167 message << data_._nested_tensor_size() << "]"; in unpack()
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/ |
| H A D | TensorProperties.cpp | 43 const auto self_nt_size = _nested_tensor_size(self); in nested_is_same_size()
44 const auto other_nt_size = _nested_tensor_size(other); in nested_is_same_size()
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/nested/cuda/ |
| H A D | NestedTensorTransformerFunctions.cpp | 379 grad_q = wrap_buffer(grad_q.view(-1), query.transpose(1,2)._nested_tensor_size()).transpose(1,2); in _scaled_dot_product_flash_attention_backward_nested()
380 grad_k = wrap_buffer(grad_k.view(-1), key.transpose(1,2)._nested_tensor_size()).transpose(1,2); in _scaled_dot_product_flash_attention_backward_nested()
381 grad_v = wrap_buffer(grad_v.view(-1), value.transpose(1,2)._nested_tensor_size()).transpose(1,2); in _scaled_dot_product_flash_attention_backward_nested()
|
| /aosp_15_r20/external/pytorch/torch/autograd/ |
| H A D | __init__.py | 78 out_shape = output._nested_tensor_size()
79 grad_shape = grad._nested_tensor_size()
|
| /aosp_15_r20/external/pytorch/benchmarks/nested/ |
| H A D | nested_bmm_bench.py | 37 nt_a_size = torch.ops.aten._nested_tensor_size(nt_a)
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/nested/ |
| H A D | NestedTensorUtils.cpp | 22 at::Tensor _nested_tensor_size(const at::Tensor& self) { in _nested_tensor_size() function
|
| /aosp_15_r20/external/pytorch/torch/multiprocessing/ |
| H A D | reductions.py | 422 rebuild_sizes_func, rebuild_sizes_args = reduce_tensor(nt._nested_tensor_size())
|
| /aosp_15_r20/external/pytorch/test/expect/ |
| H A D | HasDecompTest.test_has_decomposition.expect | 464 aten::_nested_tensor_size
465 aten::_nested_tensor_size.out
|
| /aosp_15_r20/external/pytorch/torch/ |
| H A D | _tensor.py | 410 self._nested_tensor_size(),
|
| H A D | overrides.py | 1452 Tensor._nested_tensor_size: lambda self: -1,
|
| /aosp_15_r20/external/pytorch/test/ |
| H A D | test_nestedtensor.py | 1622 size = nt_sum._nested_tensor_size()
2344 nt_sizes = nt._nested_tensor_size()
2356 nt_unsqueezed._nested_tensor_size()[:, size_idx],
2368 self.assertEqual(nt_squeezed._nested_tensor_size(), nt_sizes)
5138 self.assertEqual(nt1._nested_tensor_size(), nt2._nested_tensor_size())
|