How do I implement masking in TensorFlow eager execution?

Question

I am training a stateful RNN on variable length sequences (optional: see my previous question for more details).

I padded the sequences to a fixed length with the value -1.

The when batches are loaded, some samples will be entirely -1. (e.g. the batches are shape [batchsize, ...] and samples 1,6,8 may be entirely composed of -1's). I would like to:

1. not include these samples in the loss function calculation
2. not perform operations on them so as to speedup training.

Attempt 1:

I tried using tf.Keras.layers.Masking as in:

input = tf.Keras.layers.Masking(mask_val=-1)(input)

But, this doesnt seem to do anything. The subsequent operations are still performed, and as far as I can tell the samples are still included in the loss function. Why is this?

Attempt 2:

I tried making my own custom layer which would actually remove samples which are masked (see code below). For example the input shape goes from [batchsize, ...] to [adj_batchsize, ...] where adj_batchsize = batchsize - num_removed_samples. This works, but is extremely slow, since every time the input shape changes, the GPU memory needs to be reallocated which slows training down by a lot.

class MaskLayer(K.Model):

    def __init__(self, mask_val):
        '''
        This layer takes input tensor of [batchsize, ...] and returns tensor of shape [bs_out, ...]
        Where all samples composed entirely of mask_val are removed and bs_out is the number of remaining
        samples
        '''
        super(MaskLayer, self).__init__()
        self.mask_val = mask_val
        self.mask  = tf.keras.layers.Masking(mask_value=mask_val)

    def call(self, data, lbls):

        good_batches = tf.where(tf.math.reduce_all(self.mask.compute_mask(data), axis=range(1,data.ndim-1)))[:,0]
        data = tf.gather(data, good_batches)
        lbls = tf.gather(lbls, good_batches)
        return data, lbls

So what would be the best way to do this?

NOTE: This question was previously asked on stackoverflow.coom, but deleted due to lack of response. I think this will be a better home for it.