Layout Transformation Operators

CUDA Operators

Tensor recat_embedding_grad_output_cuda(Tensor grad_output, const std::vector<int64_t> &num_features_per_rank)

Tensor recat_embedding_grad_output_mixed_D_cuda(const Tensor &grad_output, const std::vector<int64_t> &dim_sum_per_rank)

Tensor recat_embedding_grad_output_mixed_D_batch_cuda(const Tensor &grad_output, const Tensor &dim_sum_per_rank, const Tensor &cumsum_dim_sum_per_rank)

CPU Operators

Tensor recat_embedding_grad_output_mixed_D_cpu(const Tensor &grad_output, const std::vector<int64_t> &dim_sum_per_rank)