How does a Keras mask from Embedding(mask_zero=True) reach an LSTM layer?

## The problem masking solves Sequences in a batch have different lengths, so they are padded to a common length — conventionally with token id 0. Those padded positions carry no information, and letting them influence the computation is a real accuracy bug: a recurrent layer would update its state on filler, and a pooled average would be diluted by zeros. Masking is Keras' mechanism for carrying "which timesteps are real" alongside the data. ## Where a mask comes from Two standard sources: - `keras.layers.Embedding(input_dim, output_dim, mask_zero=True)` — generates a boolean mask of shape `(batch, timesteps)`, true where the input id is not 0. The price is that index 0 becomes reserved: your vocabulary must not use it for a real token, and `input_dim` should account for the reserved slot. - `keras.layers.Masking(mask_value=0.0)` — for already-dense inputs (say precomputed features), marks a timestep as masked when the whole feature vector equals `mask_value`. You can also compute a mask yourself in a custom layer and return it from `compute_mask`. ## How it travels Keras attaches the mask to the tensor flowing out of the producing layer and threads it through the graph. Each downstream layer falls into one of three cases: 1. **Pass-through.** The layer sets `self.supports_masking = True`, meaning "my output has the same time structure as my input, so keep forwarding the mask unchanged". Most element-wise and shape-preserving layers do this. 2. **Transform.** The layer overrides `compute_mask(self, inputs, previous_mask)` and returns a new mask — or `None` when it destroys the time axis. A layer that pools over time should return `None`; a layer that subsamples timesteps must return the correspondingly subsampled mask. 3. **Consume.** The layer declares a `mask` parameter in `call(self, inputs, mask=None)` and uses it. `LSTM` and `GRU` do this: for a masked timestep they skip the update and carry the previous state forward, so the final state reflects only the real steps. A custom layer that does none of the three is the break point. By default `supports_masking` is false, and passing a mask into a layer that does not support masking is an error Keras reports explicitly rather than silently dropping the mask — which is a mercy, because a silently dropped mask is a bug you find weeks later in a metrics regression. ## Writing a masking-aware custom layer The minimum is one line in `__init__`: ``` self.supports_masking = True ``` That is correct for anything shape-preserving along the time axis — a custom activation, a residual add, a per-timestep projection. If your layer actually needs the mask (say it averages over valid timesteps only), declare it and use it: ``` def call(self, inputs, mask=None): ... ``` remembering that the mask is `(batch, timesteps)` while your tensor is usually `(batch, timesteps, features)` — expand the last axis and cast before multiplying, using `keras.ops.expand_dims` and `keras.ops.cast` so the layer stays backend-agnostic. If your layer collapses time — a global pooling — override `compute_mask` to return `None`, otherwise you hand a time-shaped mask to a layer whose output no longer has a time axis. ## Gotchas worth naming in an interview - **Index 0 is reserved** once you set `mask_zero=True`. A vocabulary that maps a real word to 0 will have that word silently ignored everywhere. - **Not every layer consumes it.** A `Dense` applied per timestep is shape-preserving and forwards the mask, but it still computes outputs at padded positions; those outputs are garbage that a later consumer must ignore. If you then flatten and feed a classifier, the padding is back in your computation. - **The loss.** If you compute a per-timestep loss, masked positions should be excluded — via `sample_weight` or a loss that respects the mask — or the model is optimized partly against padding. - **Attention layers use a different channel.** `keras.layers.MultiHeadAttention` takes an explicit `attention_mask` argument in its call, and `use_causal_mask=True` for autoregressive masking. That is a separate, explicit mechanism from the propagated sequence mask. - **Sorting/bucketing** batches by length reduces padding in the first place and is often a bigger win than perfecting the mask plumbing. ## The mental model A mask is a second, boolean tensor that rides alongside the first. Keras carries it for you as long as every layer on the path has said what to do with it — and your custom layer is, by default, the one layer that has not.