Docs/modelling layers

[mykolaskrynnyk]

Improves the documentation in layers/modeling by

• Adding a missing argument description to TokenAndPositionEmbedding.
• Aligning the usage of name argument in TransformerDecoder and TransformerEncoder layers with that in FNetEncoder layer as well as their own docstrings.

[mattdangerw]

First change looks good. Second I think we should skip.

[mattdangerw]

I think this is a place we are not fully consistent in KerasNLP, but I would say let's not do this to avoid clutter. Core Keras Dense, for example, does not do this https://github.com/keras-team/keras/blob/v3.0.5/keras/layers/core/dense.py#L33-L57

And it's not just name, it's also trainable, dtype, autocast. Would be a pain to replicate these in each layer. It would be great on the docs side to figure out how to add these to all layers across the site, but let's not clutter the code.

[mattdangerw]

If you have time, feel free to remove this from other layers where we are doing it!

[mykolaskrynnyk]

Thanks @mattdangerw. In this case, one thing we could do to increase consistency is to drop name – and other optional arguments passed explicitly to super().__init__, for that matter – from all the doctsrings and code as well as expand the definition of **kwargs. For example, changing FNetEncoder from this:

@keras_nlp_export("keras_nlp.layers.FNetEncoder")
class FNetEncoder(keras.layers.Layer):
    """FNet encoder.

    [...]

    Args:
        [...]
        name: string. The name of the layer. Defaults to `None`.
        **kwargs: other keyword arguments.

    [...]
    """

    def __init__(
        self,
        intermediate_dim,
        dropout=0,
        activation="relu",
        layer_norm_epsilon=1e-5,
        kernel_initializer="glorot_uniform",
        bias_initializer="zeros",
        name=None,
        **kwargs
    ):
        super().__init__(name=name, **kwargs)

to this:

@keras_nlp_export("keras_nlp.layers.FNetEncoder")
class FNetEncoder(keras.layers.Layer):
    """FNet encoder.

    [...]

    Args:
        [...]
        **kwargs: other keyword arguments passed to `keras.layers.Layer`, including `name`, `trainable`, 
            `dtype', `autocast` etc.

    [...]
    """

    def __init__(
        self,
        intermediate_dim,
        dropout=0,
        activation="relu",
        layer_norm_epsilon=1e-5,
        kernel_initializer="glorot_uniform",
        bias_initializer="zeros",
        **kwargs
    ):
        super().__init__(**kwargs)

That way the code will be tidier but the documentation will be clear on how users can leverage additional keywords.

What is your take?

[mattdangerw]

That looks good to me! Thank you!

[mattdangerw]

One nit. Maybe let's not even mention autocast, that's more of a power user feature.

**kwargs: other keyword arguments passed to `keras.layers.Layer`, including `name`, `trainable`, and `dtype`.

[mattdangerw]

@mykolaskrynnyk should I wait for the kwargs doc update or just pull in tie_weights docs for now and we do that on separate PR?

[mykolaskrynnyk]

@mattdangerw , I've just pushed the changes that we discussed earlier. For ReversibleEmbedding, I referenced keras.layers.Embedding in the docstring as that is the class it inherits from directly.

I think we can make similar changes to preprocessing layers. However, I can see that RandomSwap and RandomDeletion expect a specific dtype by default, so we might need to change the wording slightly. At any rate, that will be a different pull request. Cheers.

[mattdangerw]

lgtm!

[mattdangerw]

Thanks very much for the contribution!