How to handle padding during inference with a bi-directional RNN/LSTM?

[Chulabhaya]

Hey all! I am training a classifier bi-directional LSTM classifier model that is training on trajectories. During the training phase, I can use a batch of fixed-length trajectories and use zero padding and masking to account for trajectories in that batch that are of different lengths to ignore the padded data.

However, how do I ignore padded data during inference? I can set a fixed size for a trajectory, but if the trajectory ends early and part of my input data is padded with zeros, my bi-directional LSTM is going to process those padding values as legitimate values, right? Do I just ignore the padding being processed? I'm not sure how to handle this so any help would be greatly appreciated!

For context, my network is below:

class ScannedRNN(nn.Module):
    @partial(
        nn.scan,
        variable_broadcast="params",
        in_axes=0,
        out_axes=0,
        split_rngs={"params": False},
    )
    @nn.compact
    def __call__(self, carry, x):
        """Applies the module."""
        rnn_state = carry
        ins, resets = x
        hidden_size = ins.shape[-1]
        rnn_state = jnp.where(
            resets[:, np.newaxis],
            self.initialize_carry(hidden_size, *ins.shape[:-1]),
            rnn_state,
        )
        new_rnn_state, y = nn.GRUCell(hidden_size)(rnn_state, ins)
        return new_rnn_state, y

    @staticmethod
    def initialize_carry(hidden_size, *batch_size):
        # Use a dummy key since the default state init fn is just zeros.
        return nn.GRUCell(hidden_size, parent=None).initialize_carry(
            jax.random.PRNGKey(0), (*batch_size, hidden_size)
        )


class SkillDecoderNetwork(nn.Module):
    # homogenous agent for parameters sharing, assumes all agents have same obs and action dim
    skill_dim: int
    hidden_dim: int
    init_scale: float = 1.0

    @nn.compact
    def __call__(self, hidden_forward, hidden_backward, obs, dones):
        embedding = nn.Dense(
            self.hidden_dim,
            kernel_init=orthogonal(self.init_scale),
            bias_init=constant(0.0),
        )(obs)
        embedding = nn.relu(embedding)

        rnn_in_forward = (embedding, dones)
        hidden_forward, embedding_forward = ScannedRNN()(hidden_forward, rnn_in_forward)

        rnn_in_backward = (jnp.flip(embedding, axis=0), jnp.flip(dones, axis=0))
        hidden_backward, embedding_backward = ScannedRNN()(
            hidden_backward, rnn_in_backward
        )
        embedding_backward = jnp.flip(embedding_backward, axis=0)

        combined_embedding = jnp.concatenate(
            [embedding_forward, embedding_backward], -1
        )
        
        logits = nn.Dense(
            self.skill_dim,
            kernel_init=orthogonal(self.init_scale),
            bias_init=constant(0.0),
        )(combined_embedding)

        return hidden_forward, hidden_backward, logits