| import jax |
| import math |
| from typing import Any, Dict, Sequence, Union |
|
|
| import jax.numpy as jnp |
| from jax import dtypes, random |
| from jax.nn.initializers import Initializer |
| from typing import Callable |
| from flax import linen as nn |
|
|
| class FourierEmbs(nn.Module): |
| embed_scale: float |
| embed_dim: int |
| dtype: jnp.dtype = jnp.float32 |
|
|
| @nn.compact |
| def __call__(self, x): |
| kernel = self.param( |
| "kernel", jax.nn.initializers.normal(self.embed_scale, dtype=self.dtype), (x.shape[-1], self.embed_dim // 2) |
| ) |
| y = jnp.concatenate( |
| [jnp.cos(jnp.dot(x, kernel)), jnp.sin(jnp.dot(x, kernel))], axis=-1 |
| ) |
| return y |
|
|
| def _weight_fact(init_fn, mean, stddev, dtype=jnp.float32): |
| def init(key, shape): |
| key1, key2 = jax.random.split(key) |
| w = init_fn(key1, shape) |
| g = mean + nn.initializers.normal(stddev, dtype=dtype)(key2, (shape[-1],)) |
| g = jnp.exp(g) |
| v = w / g |
| return g, v |
|
|
| return init |
|
|
| class Dense(nn.Module): |
| features: int |
| kernel_init: Callable = nn.initializers.glorot_normal() |
| bias_init: Callable = nn.initializers.zeros |
| reparam : Union[None, Dict] = None |
| dtype: jnp.dtype = jnp.float32 |
|
|
| @nn.compact |
| def __call__(self, x): |
| if self.reparam is None: |
| kernel = self.param( |
| "kernel", self.kernel_init(dtype=self.dtype), (x.shape[-1], self.features) |
| ) |
| elif self.reparam["type"] == "weight_fact": |
| g, v = self.param( |
| "kernel", |
| _weight_fact( |
| self.kernel_init, |
| mean=self.reparam["mean"], |
| stddev=self.reparam["stddev"], |
| dtype=self.dtype |
| ), |
| (x.shape[-1], self.features), |
| ) |
|
|
| kernel = g * v |
|
|
| bias = self.param("bias", self.bias_init(dtype=self.dtype), (self.features,)) |
|
|
| y = jnp.dot(x, kernel) + bias |
|
|
| return y |
|
|
|
|
| def _compute_fans( |
| shape: tuple, |
| in_axis: Union[int, Sequence[int]] = -2, |
| out_axis: Union[int, Sequence[int]] = -1, |
| batch_axis: Union[int, Sequence[int]] = (), |
| ): |
| """Compute effective input and output sizes for a linear or convolutional layer. |
| |
| Axes not in in_axis, out_axis, or batch_axis are assumed to constitute the "receptive field" of |
| a convolution (kernel spatial dimensions). |
| """ |
| if len(shape) <= 1: |
| raise ValueError( |
| f"Can't compute input and output sizes of a {shape.rank}" |
| "-dimensional weights tensor. Must be at least 2D." |
| ) |
|
|
| if isinstance(in_axis, int): |
| in_size = shape[in_axis] |
| else: |
| in_size = math.prod([shape[i] for i in in_axis]) |
| if isinstance(out_axis, int): |
| out_size = shape[out_axis] |
| else: |
| out_size = math.prod([shape[i] for i in out_axis]) |
| if isinstance(batch_axis, int): |
| batch_size = shape[batch_axis] |
| else: |
| batch_size = math.prod([shape[i] for i in batch_axis]) |
| receptive_field_size = math.prod(shape) / in_size / out_size / batch_size |
| fan_in = in_size * receptive_field_size |
| fan_out = out_size * receptive_field_size |
| return fan_in, fan_out |
|
|
|
|
| def custom_uniform( |
| numerator: float = 6, |
| mode: str = "fan_in", |
| dtype: jnp.dtype = jnp.float32, |
| in_axis: Union[int, Sequence[int]] = -2, |
| out_axis: Union[int, Sequence[int]] = -1, |
| batch_axis: Sequence[int] = (), |
| distribution: str = "uniform", |
| ) -> Initializer: |
| """Builds an initializer that returns real uniformly-distributed random arrays. |
| |
| :param numerator: the numerator of the range of the random distribution. |
| :type numerator: float |
| :param mode: the mode for computing the range of the random distribution. |
| :type mode: str |
| :param dtype: optional; the initializer's default dtype. |
| :type dtype: jnp.dtype |
| :param in_axis: the axis or axes that specify the input size. |
| :type in_axis: Union[int, Sequence[int]] |
| :param out_axis: the axis or axes that specify the output size. |
| :type out_axis: Union[int, Sequence[int]] |
| :param batch_axis: the axis or axes that specify the batch size. |
| :type batch_axis: Sequence[int] |
| :param distribution: the distribution of the random distribution. |
| :type distribution: str |
| |
| :return: An initializer that returns arrays whose values are uniformly distributed in |
| the range ``[-range, range)``. |
| :rtype: Initializer |
| """ |
|
|
| def init(key: jax.random.key, shape: tuple, dtype: Any = dtype) -> Any: |
| dtype = dtypes.canonicalize_dtype(dtype) |
| fan_in, fan_out = _compute_fans(shape, in_axis, out_axis, batch_axis) |
| if mode == "fan_in": |
| denominator = fan_in |
| elif mode == "fan_out": |
| denominator = fan_out |
| elif mode == "fan_avg": |
| denominator = (fan_in + fan_out) / 2 |
| else: |
| raise ValueError(f"invalid mode for variance scaling initializer: {mode}") |
| if distribution == "uniform": |
| return random.uniform( |
| key, |
| shape, |
| dtype, |
| minval=-jnp.sqrt(numerator / denominator), |
| maxval=jnp.sqrt(numerator / denominator), |
| ) |
| elif distribution == "normal": |
| return random.normal(key, shape, dtype) * jnp.sqrt(numerator / denominator) |
| elif distribution == "uniform_squared": |
| return random.uniform( |
| key, shape, dtype, minval=-numerator / denominator, maxval=numerator / denominator |
| ) |
| else: |
| raise ValueError( |
| f"invalid distribution for variance scaling initializer: {distribution}" |
| ) |
|
|
| return init |
|
|
