| /aosp_15_r20/external/pytorch/aten/src/ATen/ |
| H A D | LegacyBatchedTensorImpl.h | 63 BatchDimsRef bdims() const { in bdims() function
128 BatchDimsRef bdims) { in createBatchDimBitset()
137 inline std::bitset<kVmapNumLevels> createVmapLevelsBitset(BatchDimsRef bdims) { in createVmapLevelsBitset()
|
| H A D | LegacyBatchedTensorImpl.cpp | 9 BatchedTensorImpl::BatchedTensorImpl(Tensor value, BatchDims bdims) in BatchedTensorImpl()
116 Tensor makeBatched(const Tensor& tensor, BatchDims bdims) { in makeBatched()
133 BatchDims bdims; in addBatchDim() local
|
| H A D | LegacyVmapTransforms.cpp | 9 static bool areBdimsAtFrontInOrder(BatchDimsRef bdims) { in areBdimsAtFrontInOrder()
21 auto bdims = batched->bdims(); in permuteBatchDimsToFront() local
92 BatchDims bdims; in computeFrontBatchDimsFromLevels() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/ |
| H A D | LegacyBatching.cpp | 23 auto bdims = batched->bdims(); in has_level() local
50 auto bdims = batched->bdims(); in remove_existing_batch_dim() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/test/ |
| H A D | legacy_vmap_test.cpp | 198 BatchDims bdims = {{/*lvl*/1, /*dim*/0}, {/*lvl*/3, /*dim*/1}}; in TEST() local
207 BatchDims bdims = {{/*lvl*/1, /*dim*/1}}; in TEST() local
217 BatchDims bdims = {{/*lvl*/1, /*dim*/1}, {/*lvl*/2,/*dim*/2}, {/*lvl*/3,/*dim*/0}}; in TEST() local
305 static void checkBatchDimsEqual(BatchDimsRef bdims, BatchDimsRef expected_bdims) { in checkBatchDimsEqual()
|
| /aosp_15_r20/external/pytorch/torch/_functorch/ |
| H A D | vmap.py | 513 def wrap_batched(args, bdims, level): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A D | triangular_solve_expander.cc | 146 int bdims = b_shape.rank(); in SolveWithInvertedDiagonalBlocks() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/functorch/ |
| H A D | BatchRulesScatterOps.cpp | 20 static bool any_has_value(ArrayRef<std::optional<int64_t>> bdims) { in any_has_value()
|
| /aosp_15_r20/external/pytorch/test/ |
| H A D | test_legacy_vmap.py | 848 def slice_inputs(inputs, bdims, i): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/tests/ |
| H A D | randomized_tests.cc | 970 std::vector<int64_t> bdims(dims.begin() + skip, dims.end()); in BroadcastableToDims() local
985 auto bdims = BroadcastableToDims(dims); in BroadcastableDims() local
|
| /aosp_15_r20/external/pytorch/test/functorch/ |
| H A D | test_vmap.py | 1265 def slice_inputs(inputs, bdims, i): argument
|