Home
last modified time | relevance | path

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

/external/pytorch/aten/src/ATen/functorch/
DBatchRulesScatterOps.cpp680 auto self_logical_rank = rankWithoutBatchDim(self, self_bdim); in scatter_batch_rule() local
687 if (self_logical_rank == 0) { in scatter_batch_rule()
699 if (self_logical_rank == 0) { in scatter_batch_rule()
712 auto self_logical_rank = rankWithoutBatchDim(self, self_bdim); in scatter_batch_rule() local
721 if (self_logical_rank == 0) { in scatter_batch_rule()
737 if (self_logical_rank == 0) { in scatter_batch_rule()
797 auto self_logical_rank = rankWithoutBatchDim(self, self_bdim); in gather_batch_rule() local
804 if (self_logical_rank == 0) { in gather_batch_rule()
918 auto self_logical_rank = rankWithoutBatchDim(self, self_bdim); in diagonal_scatter_batch_rule() local
919 dim1 = maybe_wrap_dim(dim1, self_logical_rank) + 1; in diagonal_scatter_batch_rule()
[all …]
DBatchRulesReduceOps.cpp328 auto self_logical_rank = rankWithoutBatchDim(self, self_bdim); in searchsorted_batch_rule() local
384 auto self_view_ = self_logical_rank == 0 ? self_.unsqueeze(-1) : self_.flatten(1); in searchsorted_batch_rule()
386 result = self_logical_rank == 0 ? result.squeeze(-1) : result.view(self_.sizes()); in searchsorted_batch_rule()
393 auto self_view_ = self_logical_rank == 0 ? self_.unsqueeze(-1) : self_.flatten(1); in searchsorted_batch_rule()
395 result = self_logical_rank == 0 ? result.squeeze(-1) : result.view(self_.sizes()); in searchsorted_batch_rule()
DBatchRulesLinearAlgebra.cpp72 auto self_logical_rank = rankWithoutBatchDim(self, self_bdim); in mv_batch_rule() local
74 TORCH_CHECK(self_logical_rank == 2 && other_logical_rank == 1, in mv_batch_rule()
76 "Got incorrect dims for mv(a, b). a has dim ", self_logical_rank, in mv_batch_rule()
85 auto self_logical_rank = rankWithoutBatchDim(self, self_bdim); in mm_batch_rule() local
87 TORCH_CHECK(self_logical_rank == 2 && other_logical_rank == 2, in mm_batch_rule()
89 "a has dim ", self_logical_rank, in mm_batch_rule()
100 auto self_logical_rank = rankWithoutBatchDim(self, self_bdim); in bmm_batch_rule() local
102 TORCH_CHECK(self_logical_rank == 3 && other_logical_rank == 3, in bmm_batch_rule()
104 "a has dim ", self_logical_rank, in bmm_batch_rule()
DBatchRulesBinaryOps.cpp190 const auto self_logical_rank = rankWithoutBatchDim(self, self_bdim); in masked_select_batch_rule() local
191 const auto max_logical_rank = std::max(self_logical_rank, mask.dim()); in masked_select_batch_rule()
211 const auto self_logical_rank = rankWithoutBatchDim(self, self_bdim); in masked_select_backward_batch_rule() local
212 const auto max_logical_rank = std::max(self_logical_rank, mask.dim()); in masked_select_backward_batch_rule()
DBatchRulesRandomness.cpp78 auto self_logical_rank = rankWithoutBatchDim(self_value, self_bdim); in bernoulli_inplace_Tensor_batching_rule() local
80 auto max_logical_rank = std::max(self_logical_rank, other_logical_rank); in bernoulli_inplace_Tensor_batching_rule()