Allow jagged_index_select backward to accept pre-computed output shape

[mrshenli]

Summary: Save num_dense_output_rows computed during the forward pass and use it to avoid blocking .item() call during backward.

Differential Revision: D54173841