class FlaxTimestepEmbedding(nn.Module):
    r"""
    Time step Embedding Module. Learns embeddings for input time steps.

    Args:
        time_embed_dim (`int`, *optional*, defaults to `32`):
            Time step embedding dimension.
        dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`):
            The data type for the embedding parameters.
    """

    time_embed_dim: int = 32
    dtype: jnp.dtype = jnp.float32

    @nn.compact
    def __call__(self, temb):
        temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name="linear_1")(temb)
        temb = nn.silu(temb)
        temb = nn.Dense(self.time_embed_dim, dtype=self.dtype, name="linear_2")(temb)
        return temb
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/embeddings_flax.py