Home
last modified time | relevance | path

Searched defs:wrap_dim (Results 1 – 6 of 6) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/
H A DGatedLinearUnit.cpp28 auto wrap_dim = maybe_wrap_dim(dim, self.dim()); in TORCH_META_FUNC() local
57 auto wrap_dim = maybe_wrap_dim(dim, input.dim()); in glu_backward_cpu_out() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/cpu/
H A DReduceOpsKernel.cpp80 auto wrap_dim = maybe_wrap_dim(dim, self.dim()); in cumsum_cpu_kernel() local
99 auto wrap_dim = maybe_wrap_dim(dim, self.dim()); in cumprod_cpu_kernel() local
118 auto wrap_dim = maybe_wrap_dim(dim, self.dim()); in logcumsumexp_cpu_kernel() local
H A DTensorCompareKernel.cpp173 auto wrap_dim = maybe_wrap_dim(dim, self.dim()); in aminmax_kernel() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DActivation.cpp32 auto wrap_dim = maybe_wrap_dim(dim, input.dim()); in glu_backward_cuda_out() local
H A DScanKernels.cpp63 const auto wrap_dim = maybe_wrap_dim(dim, self.dim()); in _logcumsumexp_out_cuda() local
/aosp_15_r20/external/pytorch/test/
H A Dtest_linalg.py2328 def wrap_dim(dim, ndims): function