Home
last modified time | relevance | path

Searched defs:wrapped_dim (Results 1 – 8 of 8) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/
H A DSorting.cpp226 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 DTensorShape.cpp2748 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 DTensorAdvancedIndexing.cpp161 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 DReduceOps.cpp888 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 DSoftMax.cpp40 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 DBatchRulesHelper.cpp48 auto wrapped_dim = maybe_wrap_dim(logical_dim, rank); in getPhysicalDim() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/nested/
H A DNestedTensorMath.cpp669 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 Dshape_inference.cpp1098 auto wrapped_dim = at::maybe_wrap_dim(dim, tensors[0].ndimension() + 1); in compute_shape_stack() local