Jagged Tensor Operators

Jagged Tensor solves the issue when rows in dimension are of different length. This often occurs in sparse feature inputs in recommender systems, as well as natural language processing system batched inputs.

CUDA Operators

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 where x is jagged, y is dense, and output is jagged

CPU Operators

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)

Output = x + y where x is jagged, y and output are dense

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)

Output = x + y where x is jagged, y is dense, and output is jagged

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)