Searched refs:logical_rank (Results 1 – 5 of 5) sorted by relevance
| /external/pytorch/aten/src/ATen/functorch/ |
| D | BatchRulesViews.cpp | 296 auto logical_rank = rankWithoutBatchDim(self, bdim); in roll_batch_rule() local
297 if (logical_rank == 0) { in roll_batch_rule()
312 auto logical_rank = rankWithoutBatchDim(self, self_bdim); in diagonal_batching_rule() local
314 auto dim1_ = maybe_wrap_dim(dim1, logical_rank) + 1; in diagonal_batching_rule()
315 auto dim2_ = maybe_wrap_dim(dim2, logical_rank) + 1; in diagonal_batching_rule()
323 auto logical_rank = rankWithoutBatchDim(grad_input, grad_input_bdim); in diagonal_backward_batch_rule() local
325 dim1 = maybe_wrap_dim(dim1, logical_rank + 1) + 1; in diagonal_backward_batch_rule()
326 dim2 = maybe_wrap_dim(dim2, logical_rank + 1) + 1; in diagonal_backward_batch_rule()
395 auto logical_rank = rankWithoutBatchDim(grad_input, grad_input_bdim); in select_backward_batch_rule() local
397 dim = maybe_wrap_dim(dim, logical_rank + 1) + 1; in select_backward_batch_rule()
[all …]
|
| D | BatchRulesPooling.cpp | 21 auto logical_rank = rankWithoutBatchDim(self, self_bdim); in max_pool_with_indices_batch_rule_helper() local
22 TORCH_INTERNAL_ASSERT(logical_rank == n + 1 || logical_rank == n + 2); in max_pool_with_indices_batch_rule_helper()
24 if (logical_rank == n + 1) { in max_pool_with_indices_batch_rule_helper()
|
| D | BatchRulesHelper.cpp | 71 …aybePadToLogicalRank(const Tensor& tensor, std::optional<int64_t> has_bdim, int64_t logical_rank) { in maybePadToLogicalRank() argument
76 if (tensor_logical_rank >= logical_rank) { in maybePadToLogicalRank()
80 for (int64_t i = 0; i < logical_rank - tensor_logical_rank; i++) { in maybePadToLogicalRank()
|
| D | BatchRulesHelper.h | 42 …maybePadToLogicalRank(const Tensor& tensor, std::optional<int64_t> has_bdim, int64_t logical_rank);
327 const auto logical_rank = rankWithoutBatchDim(value, bdim);
330 is_no_batch_dim_case = (logical_rank == feature_rank);
337 TORCH_INTERNAL_ASSERT(logical_rank == feature_rank);
345 TORCH_INTERNAL_ASSERT(logical_rank == feature_rank + 1);
|
| D | BatchRulesNorm.cpp | 31 static Tensor padRight(const Tensor& tensor, std::optional<int64_t> has_bdim, int64_t logical_rank)… in padRight() argument
34 if (tensor_logical_rank >= logical_rank) { in padRight()
38 for (int64_t i = 0; i < logical_rank - tensor_logical_rank; i++) { in padRight()
|