Home
last modified time | relevance | path

Searched defs:bdim_size (Results 1 – 9 of 9) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/functorch/
H A DBatchRulesModules.cpp15 static Tensor getStepTensor(const Tensor& indices, const c10::SymInt& bdim_size, const c10::SymInt&… in getStepTensor()
61 const auto bdim_size = grad.sym_size(*grad_bdim); in embedding_dense_backward_batch_rule() local
68 const auto bdim_size = indices.size(*indices_bdim); in embedding_dense_backward_batch_rule() local
167 int64_t bdim_size) { in grid_sample_backward_helper_out()
211 int64_t bdim_size = std::get<3>(new_bw_input); in cudnn_grid_sample_backward_batch_rule() local
H A DBatchRulesReduceOps.cpp343 auto bdim_size = get_bdim_size2( in searchsorted_batch_rule() local
372 auto bdim_size = self.size(*self_bdim); in searchsorted_batch_rule() local
391 auto bdim_size = buckets.size(*buckets_bdim); in searchsorted_batch_rule() local
H A DBatchRulesNorm.cpp66 std::optional<int64_t> bdim_size; in batch_norm_batch_rule() local
156 …const auto bdim_size = get_bdim_size4(grad_out, grad_out_bdim, input, input_bdim, running_mean, ru… in batch_norm_backward_no_weight_bias_batch_rule() local
312 const auto bdim_size = input_value.size(*input_bdim); in native_group_norm_plumbing() local
351 const auto bdim_size = get_bdim_size2(grad_out, grad_out_bdim, input, input_bdim); in group_norm_backward_no_weight_bias_batch_rule() local
553 const auto bdim_size = get_bdim_size2(grad_out, grad_out_bdim, input, input_bdim); in native_layer_norm_backward_no_weight_bias_batch_rule() local
H A DBatchRulesLoss.cpp107 const auto bdim_size = get_bdim_size2(self_value, self_bdim, target_value, target_bdim); in binary_cross_entropy_plumbing() local
144 const auto bdim_size = get_bdim_size3( in binary_cross_entropy_backward_plumbing() local
H A DLegacyVmapTransforms.cpp122 c10::SymInt bdim_size = -1; in logicalToPhysical() local
174 auto bdim_size = -1; in logicalToPhysical() local
H A DBatchRulesFactory.cpp22 const auto bdim_size = tensor.sym_size(batch_dim.value()); in apply() local
41 const auto bdim_size = tensor.size(batch_dim.value()); in apply() local
H A DBatchRulesPooling.cpp31 auto bdim_size = self.size(*self_bdim); in max_pool_with_indices_batch_rule_helper() local
H A DBatchRulesLinearAlgebra.cpp279 auto bdim_size = get_bdim_size2(LU, LU_bdim, pivots, pivots_bdim); in linalg_lu_unpack_batch_rule() local
310 auto bdim_size = get_bdim_size2(LU, LU_bdim, pivots, pivots_bdim); in linalg_lu_solve_batch_rule() local
H A DLegacyBatchingRegistrations.cpp539 std::optional<int64_t> bdim_size = std::nullopt; in cat_batching_rule() local