[Docs] Add docstrings for nn.Module classes and core APIs in relax.frontend.nn#19387
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request significantly improves the documentation for the Relax frontend's neural network modules and core functions. It adds detailed docstrings, including parameter descriptions and return types, for various layers such as Linear, Convolutional, and Normalization modules, as well as the JIT compilation and KVCache utilities. The review feedback suggests a minor improvement to the TimestepEmbedding documentation to explicitly mention the supported activation functions, ensuring consistency with other documented parameters.
| Dimensionality of the output. If ``None``, defaults to *time_embed_dim*. | ||
|
|
||
| post_act_fn : Optional[str] | ||
| Optional post-activation applied after the second linear layer. |
There was a problem hiding this comment.
The documentation for post_act_fn could be more specific. Similar to act_fn, it seems only "silu" is supported for post-activation. It would be helpful to mention this limitation in the docstring to avoid runtime errors.
| Optional post-activation applied after the second linear layer. | |
| Optional post-activation applied after the second linear layer. Currently only ``"silu"`` is supported. |
2 participants
This pr adds some docstrings for important APIs