Pooled Embeddings Operators

This section includes CUDA and CPU operators for various operations with pooled embeddings, including merge and permutation operators.

Merge Operators

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

Permutation Operators

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_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_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)