| 318 | |
| 319 | |
| 320 | class RMSNorm(nn.Module): |
| 321 | dim: int |
| 322 | eps: float=1e-6 |
| 323 | dtype: jnp.dtype=jnp.float32 |
| 324 | param_dtype: jnp.dtype=jnp.float32 |
| 325 | |
| 326 | def setup(self) -> None: |
| 327 | self.weight = self.param( |
| 328 | 'kernel', |
| 329 | nn.initializers.ones, |
| 330 | (self.dim,), |
| 331 | self.param_dtype, |
| 332 | ) |
| 333 | |
| 334 | def _norm(self, x: jnp.ndarray) -> jnp.ndarray: |
| 335 | return x * jax.lax.rsqrt(jnp.square(x).mean(-1, keepdims=True) + self.eps) |
| 336 | |
| 337 | def __call__(self, x: jnp.ndarray) -> jnp.ndarray: |
| 338 | x = x.astype(jnp.promote_types(self.dtype, jnp.float32)) |
| 339 | output = self._norm(x).astype(self.dtype) |
| 340 | weight = jnp.asarray(self.weight, self.dtype) |
| 341 | return output * weight |
| 342 | |
| 343 | |
| 344 | def precompute_freqs_cis(dim: int, max_position_embedding: int, theta: float=10000.0, dtype: jnp.dtype=jnp.float32) -> jnp.ndarray: |