jcm/models/layerspp.py [33:46]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
class GaussianFourierProjection(nn.Module):
    """Gaussian Fourier embeddings for noise levels."""

    embedding_size: int = 256
    scale: float = 1.0

    @nn.compact
    def __call__(self, x):
        W = self.param(
            "W", jax.nn.initializers.normal(stddev=self.scale), (self.embedding_size,)
        )
        W = jax.lax.stop_gradient(W)
        x_proj = x[:, None] * W[None, :] * 2 * jnp.pi
        return jnp.concatenate([jnp.sin(x_proj), jnp.cos(x_proj)], axis=-1)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



jcm/models/wideresnet_noise_conditional.py [213:226]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
class GaussianFourierProjection(nn.Module):
    """Gaussian Fourier embeddings for noise levels."""

    embedding_size: int = 256
    scale: float = 1.0

    @nn.compact
    def __call__(self, x):
        W = self.param(
            "W", jax.nn.initializers.normal(stddev=self.scale), (self.embedding_size,)
        )
        W = jax.lax.stop_gradient(W)
        x_proj = x[:, None] * W[None, :] * 2 * jnp.pi
        return jnp.concatenate([jnp.sin(x_proj), jnp.cos(x_proj)], axis=-1)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -