Home
last modified time | relevance | path

Searched refs:output_shards_chunks (Results 1 – 2 of 2) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/functorch/
H A DBatchedFallback.cpp378 auto output_shards_chunks = MatrixRef<Tensor>(output_shards, num_batches); in batchedTensorForLoopFallback() local
380 auto shards = output_shards_chunks[return_idx]; in batchedTensorForLoopFallback()
502 auto output_shards_chunks = MatrixRef<Tensor>(output_shards, num_components); in batchedNestedTensorForLoopFallback() local
504 auto shards = output_shards_chunks[return_idx]; in batchedNestedTensorForLoopFallback()
/aosp_15_r20/external/pytorch/aten/src/ATen/
H A DLegacyBatchedFallback.cpp358 auto output_shards_chunks = MatrixRef<Tensor>(output_shards, num_batches); in batchedTensorForLoopFallback() local
360 auto shards = output_shards_chunks[return_idx]; in batchedTensorForLoopFallback()