google / google/flax

Add a `mask` argument to `RNN.__call__`

Open
#3,507 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

Some modules that receive a `mask` argument in their `__call__` method include:

* [BatchNorm](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.BatchNorm)
* [SelfAttention](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.SelfAttention.__call__)
* [MultiHeadDotProductAttention](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.MultiHeadDotProductAttention)

This is useful because it allows one to dynamically exclude elements from a reduction.

**Feature request:** Add a `mask` argument to [RNN](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.RNN.html#flax.linen.RNN.__call__)'s `__call__` method that makes the RNN's carry skip over masked-out elements. Thus the RNN behaves as if the masked-out elements aren't there. This can be implemented by using a _masked scan_ function. An example implementation of the latter can be found at https://github.com/google/jax/issues/18664.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.