Home
last modified time | relevance | path

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

/aosp_15_r20/external/pytorch/aten/src/ATen/functorch/
H A DBatchRulesReduceOps.cpp319 std::optional<int64_t> sorted_sequence_bdim, in searchsorted_batch_rule() argument
327 auto buckets_logical_rank = rankWithoutBatchDim(sorted_sequence, sorted_sequence_bdim); in searchsorted_batch_rule()
333 auto buckets = moveBatchDimToFront(sorted_sequence, sorted_sequence_bdim); in searchsorted_batch_rule()
335 if (sorted_sequence_bdim.has_value()) { in searchsorted_batch_rule()
342 if (sorted_sequence_bdim.has_value() != sorter_bdim.has_value()) { in searchsorted_batch_rule()
344 sorted_sequence, sorted_sequence_bdim, in searchsorted_batch_rule()
347 buckets = ensure_has_bdim(buckets, sorted_sequence_bdim.has_value(), bdim_size); in searchsorted_batch_rule()