Searched refs:BatchedNestedTensor (Results 1 – 9 of 9) sorted by relevance
| /aosp_15_r20/external/pytorch/aten/src/ATen/functorch/ |
| H A D | BatchedFallback.cpp | 411 c10::impl::ExcludeDispatchKeyGuard nt_guard(DispatchKey::BatchedNestedTensor); in batchedNestedTensorForLoopFallback()
457 c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::BatchedNestedTensor); in batchedNestedTensorForLoopFallback()
487 c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::BatchedNestedTensor); in batchedNestedTensorForLoopFallback()
505 c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::BatchedNestedTensor); in batchedNestedTensorForLoopFallback()
|
| H A D | BatchedTensorImpl.cpp | 18 value.is_nested() ? DispatchKey::BatchedNestedTensor : DispatchKey::FuncTorchBatched), in BatchedTensorImpl()
27 if (value_.is_nested() || value_.key_set().has(DispatchKey::BatchedNestedTensor)) { in BatchedTensorImpl()
|
| H A D | LegacyBatchingRegistrations.cpp | 704 c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::BatchedNestedTensor); in nested_cat_batching_rule()
758 TORCH_LIBRARY_IMPL(_, BatchedNestedTensor, m) { in TORCH_LIBRARY_IMPL() argument
763 TORCH_LIBRARY_IMPL(aten, BatchedNestedTensor, m) { in TORCH_LIBRARY_IMPL() argument
|
| H A D | BatchedTensorImpl.h | 116 tensor.unsafeGetTensorImpl()->key_set().has(DispatchKey::BatchedNestedTensor); in isBatchedTensor()
|
| H A D | Interpreter.cpp | 41 return DispatchKeySet({DispatchKey::FuncTorchBatched, DispatchKey::BatchedNestedTensor}); in keysForEnteringDynamicLayer()
|
| /aosp_15_r20/external/pytorch/c10/core/ |
| H A D | DispatchKey.cpp | 157 case DispatchKey::BatchedNestedTensor: in toString()
304 {"BatchedNestedTensor", c10::DispatchKey::BatchedNestedTensor}, in parseDispatchKey()
|
| H A D | DispatchKey.h | 374 BatchedNestedTensor, enumerator
|
| /aosp_15_r20/external/pytorch/torch/csrc/functorch/ |
| H A D | init.cpp | 149 out_dim == 0 || !self.key_set().has(DispatchKey::BatchedNestedTensor), in _remove_batch_dim()
|
| /aosp_15_r20/external/pytorch/torchgen/ |
| H A D | model.py | 119 BatchedNestedTensor = auto() variable in DispatchKey
|