[Question][Jagged Tensor] Is there concat ops for jagged tensor?

[JacoCheung]

Hi experts, I would like to know is there a way to concat 2 jagged tensors along the sparse dim?
For example

j1 = jagged_tensor.JaggedTensor(
    values=torch.tensor([1, 2, 3, 4]), lengths=torch.tensor([3, 1])
)

j2 = jagged_tensor.JaggedTensor(
    values=torch.tensor([5, 6, 7]), lengths=torch.tensor([2, 1])
)

j3 = torch.ops.fbgemm.jagged_concat(j1, j2) # the concat op
# expected result:
expected = jagged_tensor.JaggedTensor(
    values=torch.tensor([1, 2, 3, 5, 6, 4, 7]), lengths=torch.tensor([5, 2])
)

[JacoCheung]

I found out relevant op by Generative-Recommender.
Wondering if FBGEMM team could take it ?

Another request is concat more than 2 tensors that share the same batchsize like torch.cat which accepts a list of tensors.

[JacoCheung]

Hi fbgemm team. Our teammates actually have some progress on this. Will file a PR in the near future.