Mask Select with Variable Lengths

Hey guys,

I was wondering, is there a way to mask select a Tensor that would end up with inconsistent dimensions? For example: we have tensor [[1,2,], [3,4]] and mask [[1,0],[1,1]]. When we put it through a custom_mask_function:

custom_mask_function([[1,2,], [3,4]], [[1,0],[1,1]]) = [[1], [3,4]]

As you can see, the tensor does not have consistent dimensions as the first array is of length one and the second one of length two. Is creating something like custom_mask_function possible?

Thanks in advance!