Jagged Tensor 算子

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

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)

output = 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)