锯齿状张量算子

锯齿状张量解决了维度中行长不同的问题。这通常发生在推荐系统中的稀疏特征输入以及自然语言处理系统的批处理输入中。

CUDA 算子

at::Tensor jagged_to_padded_dense_forward(const Tensor &values, const std::vector<Tensor> &offsets, c10::SymIntArrayRef max_lengths, const double padding_value)

std::tuple<Tensor, std::vector<Tensor>> jagged_dense_elementwise_add_jagged_output_cuda(const Tensor &x_values, const std::vector<Tensor> &x_offsets, const Tensor &y)

输出 = x + y，其中 x 为锯齿状，y 为密集型，输出为锯齿状

CPU 算子

Tensor jagged_to_padded_dense(const Tensor &values, const std::vector<Tensor> &offsets, const c10::SymIntArrayRef max_lengths, const double padding_value)

Tensor jagged_dense_elementwise_add(const Tensor &x_values, const std::vector<Tensor> &x_offsets, const Tensor &y)

输出 = x + y，其中 x 为锯齿状，y 和输出为密集型

std::tuple<Tensor, std::vector<Tensor>> jagged_dense_elementwise_mul(const Tensor &x_values, const std::vector<Tensor> &x_offsets, const Tensor &y)

Tensor batched_dense_vec_jagged_2d_mul(const Tensor &v, const Tensor &a_values, const Tensor &a_offsets)

std::tuple<Tensor, std::vector<Tensor>> dense_to_jagged(const Tensor &dense, const std::vector<Tensor> &offsets, std::optional<at::SymInt> total_L)

std::tuple<Tensor, std::vector<Tensor>> jagged_dense_elementwise_add_jagged_output(const Tensor &x_values, const std::vector<Tensor> &x_offsets, const Tensor &y)

输出 = x + y，其中 x 为锯齿状，y 为密集型，输出为锯齿状

Tensor jagged_1d_to_dense(Tensor values, Tensor offsets, c10::SymInt max_L, int64_t padding_value)

Tensor jagged_2d_to_dense(Tensor values, Tensor offsets, c10::SymInt max_sequence_length)