# kd.nn

[[[Source]]](https://github.com/google-research/kauldron/tree/main/kauldron/modules/__init__.py)

```{eval-rst}
.. automodule:: kauldron.modules
  :no-members:
```

## Symbols


### Class

|  |  |
--- | ---
[kd.nn.AddEmbedding](AddEmbedding) | Helper Module for adding a PositionEmbedding e.g. in a `knn.Sequential`.
[kd.nn.AddLearnedEmbedding](AddLearnedEmbedding) | Adds learned positional embeddings to the inputs.
[kd.nn.AttentionModule](AttentionModule) | Interface specification for Attention modules.
[kd.nn.Dropout](Dropout) | Wrapper around `nn.Dropout` but using `kd.nn.train_property`.
[kd.nn.DummyModel](DummyModel) | Empty model that ignores inputs and always produces a single logit of 42.
[kd.nn.ExternalModule](ExternalModule) | Module that is defined outside Kauldron.
[kd.nn.FlatAutoencoder](FlatAutoencoder) | Very simple auto-encoder class to showcase using keys and submodules.
[kd.nn.FourierEmbedding](FourierEmbedding) | Apply Fourier position embedding to a grid of coordinates.
[kd.nn.Identity](Identity) | Module that applies the identity function to a single tensor.
[kd.nn.ImageTokenizer](ImageTokenizer) | Interface for modules that convert images into tokens.
[kd.nn.ImprovedMultiHeadDotProductAttention](ImprovedMultiHeadDotProductAttention) | Multi-head dot-product attention.
[kd.nn.LearnedEmbedding](LearnedEmbedding) | Learned positional embeddings.
[kd.nn.MultiHeadDotProductAttention](MultiHeadDotProductAttention) | Wrapper around `nn.MultiHeadDotProductAttention` using `knn.train_property`.
[kd.nn.NormModule](NormModule) | Interface specification for norm modules (to be used as type annotation).
[kd.nn.ParallelAttentionBlock](ParallelAttentionBlock) | Parallel self attention (see Vit22B paper: arxiv.org/abs/2302.05442).
[kd.nn.Patchify](Patchify) | Patchify an image, as in ViT (without linear embedding).
[kd.nn.PatchifyEmbed](PatchifyEmbed) | Patchify and linearly embed and image, as in ViT.
[kd.nn.PostNormBlock](PostNormBlock) | Post-LN Transformer layer (not recommended).
[kd.nn.PreNormBlock](PreNormBlock) | Pre-LN Transformer layer (default transformer layer).
[kd.nn.Rearrange](Rearrange) | Wrapper around `einops.rearrange` for usage e.g. in `nn.Sequential`.
[kd.nn.Reduce](Reduce) | Wrapper around `einops.reduce` for usage e.g. in `nn.Sequential`.
[kd.nn.Sequential](Sequential) | Like nn.Sequential but allows configuring input and output keys.
[kd.nn.TransformerBlock](TransformerBlock) | Interface definition for transformer blocks (for use in type annotations).
[kd.nn.TransformerMLP](TransformerMLP) | Simple MLP with a single hidden layer for use in Transformer blocks.
[kd.nn.Vit](Vit) | Basic Vision Transformer classifer with GAP.
[kd.nn.VitEncoder](VitEncoder) | Basic Vit Encoder.
[kd.nn.WrapperModule](WrapperModule) | Base class to wrapper a module.
[kd.nn.ZeroEmbedding](ZeroEmbedding) | Embedding that returns zero (for deactivating position embeddings).

### Function

|  |  |
--- | ---
[kd.nn.convert_to_fourier_features](convert_to_fourier_features) | Convert inputs to Fourier features, e.g. for positional encoding.
[kd.nn.interms_property](interms_property) | `interms` property that makes storing intermediates more convenient.
[kd.nn.set_train_property](set_train_property) | Set the `self.is_training` state to the given value.
[kd.nn.train_property](train_property) | `is_training` property.

```{toctree}
:hidden:

AddEmbedding
AddLearnedEmbedding
AttentionModule
Dropout
DummyModel
ExternalModule
FlatAutoencoder
FourierEmbedding
Identity
ImageTokenizer
ImprovedMultiHeadDotProductAttention
LearnedEmbedding
MultiHeadDotProductAttention
NormModule
ParallelAttentionBlock
Patchify
PatchifyEmbed
PostNormBlock
PreNormBlock
Rearrange
Reduce
Sequential
TransformerBlock
TransformerMLP
Vit
VitEncoder
WrapperModule
ZeroEmbedding
convert_to_fourier_features
interms_property
set_train_property
train_property
```