and works in the DTensor training use case.
    *   Improved masking support for
        [`tf.keras.layers.MultiHeadAttention`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention).
        *   Implicit masks for `query`, `key` and `value` inputs will
            automatically be used to compute a correct attention mask for the
            layer. These padding masks will be combined with any