| /external/pytorch/aten/src/ATen/functorch/ |
| D | BatchRulesModules.cpp | 36 result = reshape_dim_outof(-1, batch_size, result); in embedding_batch_rule()
75 result = reshape_dim_outof(0, bdim_size, result); in embedding_dense_backward_batch_rule()
118 out = reshape_dim_outof(1, input.sizes()[*input_bdim], out); in grid_sample_batch_rule()
124 out = reshape_dim_outof(2, grid.sizes()[*grid_bdim], out); in grid_sample_batch_rule()
130 out = reshape_dim_outof(0, input.sizes()[*grid_bdim], out); in grid_sample_batch_rule()
170 grad_input = reshape_dim_outof(*grad_input_out_bdim, bdim_size, grad_input); in grid_sample_backward_helper_out()
171 grad_grid = reshape_dim_outof(*grad_grid_out_bdim, bdim_size, grad_grid); in grid_sample_backward_helper_out()
|
| D | BatchRulesPooling.cpp | 36 reshape_dim_outof(0, bdim_size, std::get<0>(result)), 0, in max_pool_with_indices_batch_rule_helper()
37 reshape_dim_outof(0, bdim_size, std::get<1>(result)), 0); in max_pool_with_indices_batch_rule_helper()
|
| D | BatchRulesConvolution.cpp | 185 const auto grad_input = reshape_dim_outof(1, batch_size, std::get<0>(result)); in convolution_backward_input_batch_rule()
196 const auto grad_input = reshape_dim_outof(0, batch_size, std::get<0>(result)); in convolution_backward_input_batch_rule()
209 const auto grad_input = reshape_dim_outof(1, batch_size, std::get<0>(result)); in convolution_backward_input_batch_rule()
414 reshape_dim_outof(1, batch_size, std::get<0>(result)) : Tensor(); in convolution_backward_plumbing()
416 reshape_dim_outof(0, batch_size, std::get<1>(result)) : Tensor(); in convolution_backward_plumbing()
|
| D | BatchRulesLinearAlgebra.cpp | 518 const auto res0_ = reshape_dim_outof(0, batch_size, res0); in _scaled_dot_product_flash_attention_batch_rule()
519 const auto res1_ = reshape_dim_outof(0, batch_size, res1); in _scaled_dot_product_flash_attention_batch_rule()
523 const auto res8_ = return_debug_mask ? reshape_dim_outof(0, batch_size, res8) : res8; in _scaled_dot_product_flash_attention_batch_rule()
571 const auto res0_ = reshape_dim_outof(0, batch_size, res0); in _scaled_dot_product_efficient_attention_batch_rule()
572 const auto res1_ = reshape_dim_outof(0, batch_size, res1); in _scaled_dot_product_efficient_attention_batch_rule()
614 const auto res0_ = reshape_dim_outof(0, batch_size, res0); in _scaled_dot_product_cudnn_attention_batch_rule()
618 res1_ = reshape_dim_outof(0, batch_size, res1); in _scaled_dot_product_cudnn_attention_batch_rule()
627 const auto res8_ = return_debug_mask ? reshape_dim_outof(0, batch_size, res8) : res8; in _scaled_dot_product_cudnn_attention_batch_rule()
|
| D | BatchRulesNorm.cpp | 100 result0 = reshape_dim_outof(0, bdim_size.value(), result0); // [B0, C, B, *] in batch_norm_batch_rule()
102 mean = reshape_dim_outof(0, bdim_size.value(), mean); // [B0, C] in batch_norm_batch_rule()
104 rstd = reshape_dim_outof(0, bdim_size.value(), rstd); // [B0, C] in batch_norm_batch_rule()
193 result0 = reshape_dim_outof(1, bdim_size, result0); // [B, B0, C, *] in batch_norm_backward_no_weight_bias_batch_rule()
316 result0 = makeBatched(reshape_dim_outof(0, bdim_size, std::get<0>(result)), 0, cur_level); in native_group_norm_plumbing()
317 mean = makeBatched(reshape_dim_outof(0, bdim_size, std::get<1>(result)), 0, cur_level); in native_group_norm_plumbing()
318 rstd = makeBatched(reshape_dim_outof(0, bdim_size, std::get<2>(result)), 0, cur_level); in native_group_norm_plumbing()
369 result0 = reshape_dim_outof(0, bdim_size, result0); in group_norm_backward_no_weight_bias_batch_rule()
412 const auto reshaped_input = reshape_dim_outof(1, group, input); in native_group_norm_backward_plumbing()
|
| D | BatchRulesHelper.h | 29 TORCH_API Tensor reshape_dim_outof(int64_t src, int64_t size1, const Tensor& x);
283 (*stack)[idx] = makeBatched(reshape_dim_outof(0, batch_size, ret.toTensor()), 0, cur_level);
362 (*stack)[idx] = makeBatched(reshape_dim_outof(0, batch_size, ret.toTensor()), 0, cur_level);
|
| D | BatchRulesHelper.cpp | 114 Tensor reshape_dim_outof(int64_t src, int64_t size1, const Tensor& x) { in reshape_dim_outof() function
|
| D | BatchRulesReduceOps.cpp | 375 result = reshape_dim_outof(-1, bdim_size, result); in searchsorted_batch_rule()
|
| D | BatchRulesScatterOps.cpp | 1073 self_ = reshape_dim_outof(dim, batch_size, self_); in index_fill_batch_rule_helper()
|
| /external/pytorch/torch/csrc/functorch/ |
| D | init.cpp | 508 m.def("reshape_dim_outof", &at::functorch::reshape_dim_outof); in initFuncTorchBindings()
|
| /external/pytorch/test/functorch/ |
| D | test_vmap.py | 44 from torch._C._functorch import reshape_dim_into, reshape_dim_outof
1107 y = reshape_dim_outof(0, 2, x)
1110 y = reshape_dim_outof(1, 4, x)
1113 y = reshape_dim_outof(2, 6, x)
1116 y = reshape_dim_outof(-1, 6, x)
1121 y = reshape_dim_outof(-1, 6, x)
|