Loss on masked tensors

Suppose I have logits like

[[4.3, -0.5, -2.7, 0, 0], [0.5, 2.3, 0, 0, 0]]

where clearly the last two in the first example and last three in the second example are masked and should not affect loss and gradient computations. How to compute cross-entropy loss between this logits and corresponding labels ? For sanity, the labels for this example can be something like

[[1, 0, 0, 0, 0], [0, 1, 0, 0, 0]]

(One issue: Softmax, followed by log, on the logits will be applicable for the masked zeroes too and tf's cross entropy method will consider the loss for those elements too.)

(Also, you can think about the problem like this: I have logits of varying lengths in a batch, i.e. my logits were length 3 and 2 for eg.1 and eg.2 respectively. Same is followed by the labels.)

from Loss on masked tensors