Home
last modified time | relevance | path

Searched defs:indices_arg (Results 1 – 11 of 11) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/
H A DEmbedding.cpp40 auto indices_arg = TensorArg(indices, "indices", 1); in embedding_symint() local
79 auto indices_arg = TensorArg(indices_, "indices", 2); in embedding_sparse_backward() local
116 auto indices_arg = TensorArg(indices, "indices", 2); in embedding_dense_backward_cpu() local
184 auto indices_arg = TensorArg(indices, "indices", 2); in embedding_renorm_cpu_() local
H A DEmbeddingBag.cpp877 auto indices_arg = TensorArg(indices, "indices", 1); in check_arguments() local
1405 auto indices_arg = TensorArg(indices, "indices", 1); in _embedding_bag_backward_symint() local
1671 auto indices_arg = TensorArg(indices, "indices", 1); in _embedding_bag_per_sample_weights_backward_cpu_template() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DMaxUnpooling.cu132 indices_arg{indices_, "indices_", 3}; in max_unpooling2d_forward_out_cuda() local
310 indices_arg{indices_, "indices_", 3}; in max_unpooling3d_forward_out_cuda() local
413 self_arg{self_, "self_", 3}, indices_arg{indices_, "indices_", 4}; in max_unpooling2d_backward_out_cuda() local
521 TensorArg self_arg{self_, "self_", 1}, indices_arg{indices_, "indices_", 2}, in max_unpooling3d_backward_out_cuda() local
H A DScanKernels.cpp30 TensorArg indices_arg{ indices, "indices", 2 }; in cummax_helper_cuda() local
47 TensorArg indices_arg{ indices, "indices", 2 }; in cummin_helper_cuda() local
H A DEmbedding.cu252 auto indices_arg = TensorArg(indices_, "indices", 1); in embedding_dense_backward_cuda() local
350 auto indices_arg = TensorArg(indices, "indices", 1); in embedding_renorm_cuda_() local
H A DAdaptiveMaxPooling2d.cu210 TensorArg indices_arg{indices, "indices", 2}; in TORCH_IMPL_FUNC() local
334 TensorArg indices_arg{indices, "indices", 4}; in TORCH_IMPL_FUNC() local
H A DAdaptiveMaxPooling3d.cu311 TensorArg indices_arg{indices, "indices", 2}; in TORCH_IMPL_FUNC() local
393 TensorArg indices_arg{indices, "indices", 4}; in TORCH_IMPL_FUNC() local
H A DTensorTopK.cpp52 …TensorArg topK_arg{values, "topK", 1}, indices_arg{indices, "indices", 2}, input_arg{self, "self",… in TORCH_IMPL_FUNC() local
H A DEmbeddingBag.cu354 auto indices_arg = TensorArg(indices, "indices", 1); in _embedding_bag_cuda() local
445 auto indices_arg = TensorArg(indices, "indices", 1); in _embedding_bag_dense_backward_cuda() local
H A DDilatedMaxPool3d.cu294 TensorArg indices_arg{ indices, "indices", 2 }; in max_pool3d_with_indices_out_cuda_template() local
425 TensorArg indices_arg{ indices, "indices", 4 }; in max_pool3d_with_indices_backward_out_cuda_template() local
H A DDilatedMaxPool2d.cu306 TensorArg indices_arg{ indices, "indices", 2 }; in TORCH_IMPL_FUNC() local
432 TensorArg indices_arg{ indices_, "indices", 4 }; in TORCH_IMPL_FUNC() local