model/init.py

from typing import Callable

import jax
import jax.numpy as jnp
from jaxtyping import Array


def uniform_between(a: float, b: float, dtype=jnp.float32) -> Callable:
    def init(key, shape, dtype=dtype) -> Array:
        return jax.random.uniform(key, shape, dtype=dtype, minval=a, maxval=b)
    return init


def linear_up(scale: float) -> Callable:
    def init(key, shape, dtype=jnp.float32) -> Array:
        assert shape[-2] == 2
        keys = jax.random.split(key, 2)
        norm = jnp.pi * scale * (
                jax.random.uniform(keys[0], shape=(1, shape[-1])) ** .5)
        theta = 2 * jnp.pi * jax.random.uniform(keys[1], shape=(1, shape[-1]))
        x = norm * jnp.cos(theta)
        y = norm * jnp.sin(theta)
        return jnp.concatenate([x, y], axis=-2).astype(dtype)
    return init