Searched defs:wrapped_dim (Results 1 – 8 of 8) sorted by relevance
/aosp_15_r20/external/pytorch/aten/src/ATen/native/ |
H A D | Sorting.cpp | 226 int64_t wrapped_dim) { in quantile_output_shape() 253 int64_t wrapped_dim, in quantile_compute() 387 int64_t wrapped_dim = at::maybe_wrap_dim(original_dim.value_or(0), self.dim()); in quantile_out_impl() local 406 int64_t wrapped_dim = at::maybe_wrap_dim(original_dim.value_or(0), self.dim()); in quantile_impl() local
|
H A D | TensorShape.cpp | 2748 auto wrapped_dim = at::native::preprocess_chunk_cat_inputs(tensors, dim, num_chunks); in _chunk_cat() local 2753 auto wrapped_dim = at::native::preprocess_chunk_cat_inputs(tensors, dim, num_chunks); in _chunk_cat_out() local 2762 auto wrapped_dim = maybe_wrap_dim(dim, tensors[0].ndimension()+1); in stack() local 2792 auto wrapped_dim = maybe_wrap_dim(dim, tensors[0].ndimension()+1); in stack_out() local
|
H A D | TensorAdvancedIndexing.cpp | 161 int64_t wrapped_dim = at::maybe_wrap_dim(dim, self.dim()); in TORCH_META_FUNC() local 195 int64_t wrapped_dim = at::maybe_wrap_dim(dim, self.dim()); in scatter_meta_impl() local
|
H A D | ReduceOps.cpp | 888 int64_t wrapped_dim = maybe_wrap_dim(dim, self.dim(), false); in diff_check_compatible_shape() local
|
/aosp_15_r20/external/pytorch/aten/src/ATen/native/mkldnn/ |
H A D | SoftMax.cpp | 40 const int64_t wrapped_dim = maybe_wrap_dim(dim, self.dim()); in mkldnn_softmax() local
|
/aosp_15_r20/external/pytorch/aten/src/ATen/functorch/ |
H A D | BatchRulesHelper.cpp | 48 auto wrapped_dim = maybe_wrap_dim(logical_dim, rank); in getPhysicalDim() local
|
/aosp_15_r20/external/pytorch/aten/src/ATen/native/nested/ |
H A D | NestedTensorMath.cpp | 669 int64_t wrapped_dim = at::maybe_wrap_dim(dim, ndim + 1); in unsqueeze_nested() local
|
/aosp_15_r20/external/pytorch/torch/csrc/lazy/core/ |
H A D | shape_inference.cpp | 1098 auto wrapped_dim = at::maybe_wrap_dim(dim, tensors[0].ndimension() + 1); in compute_shape_stack() local
|