Pooled Embeddings 運算子

本節包含 CUDA 和 CPU 運算子，用於各種 pooled embeddings 運算，包括合併和置換運算子。

合併運算子

std::vector<at::Tensor> all_to_one_device(std::vector<at::Tensor> inputTensors, at::Device target_device)

置換運算子

at::Tensor permute_pooled_embs_split_gpu(const at::Tensor &pooled_embs, const at::Tensor &offset_dim_list, const at::Tensor &permute_list, const at::Tensor &inv_offset_dim_list, const at::Tensor &inv_permute_list)

at::Tensor permute_pooled_embs_auto_grad_split_gpu(const at::Tensor &pooled_embs, const at::Tensor &offset_dim_list, const at::Tensor &permute_list, const at::Tensor &inv_offset_dim_list, const at::Tensor &inv_permute_list)

Tensor permute_pooled_embs_auto_grad_gpu(const Tensor &pooled_embs, const Tensor &offset_dim_list, const Tensor &permute_list, const Tensor &inv_offset_dim_list, const Tensor &inv_permute_list)

at::Tensor permute_pooled_embs_cpu_impl(const at::Tensor &pooled_embs, const at::Tensor &offset_dim_list, const at::Tensor &permute_list, const at::Tensor &inv_offset_dim_list, const at::Tensor &inv_permute_list, const bool &allow_duplicates)

at::Tensor permute_pooled_embs_split_cpu(const at::Tensor &pooled_embs, const at::Tensor &offset_dim_list, const at::Tensor &permute_list, const at::Tensor &inv_offset_dim_list, const at::Tensor &inv_permute_list)

at::Tensor permute_pooled_embs_auto_grad_split_cpu(const at::Tensor &pooled_embs, const at::Tensor &offset_dim_list, const at::Tensor &permute_list, const at::Tensor &inv_offset_dim_list, const at::Tensor &inv_permute_list)

at::Tensor permute_pooled_embs_auto_grad(const Tensor &pooled_embs, const Tensor &offset_dim_list, const Tensor &permute_list, const Tensor &inv_offset_dim_list, const Tensor &inv_permute_list)

at::Tensor permute_pooled_embs_auto_grad_cpu(const Tensor &pooled_embs, const Tensor &offset_dim_list, const Tensor &permute_list, const Tensor &inv_offset_dim_list, const Tensor &inv_permute_list)