<commit_before>from typing import Callable

from .config import registry
from .types import Array
from .util import get_array_module, copy_array, partial


def xavier_uniform_init(data: Array, *, inplace: bool = False) -> Array:
    xp = get_array_module(data)
    scale = xp.sqrt(6.0 / (data.shape[0] + data.shape[1]))
    if inplace:
        copy_array(data, xp.random.uniform(-scale, scale, data.shape))
        return data
    else:
        return xp.random.uniform(-scale, scale, data.shape)


@registry.initializers("xavier_uniform_init.v0")
def configure_xavier_uniform_init(*, inplace: bool = False) -> Callable[[Array], Array]:
    return partial(xavier_uniform_init, inplace=inplace)


def zero_init(data: Array, *, inplace: bool = False) -> Array:
    if inplace:
        data.fill(0.0)
        return data
    else:
        xp = get_array_module(data)
        return xp.zeros_like(data)


@registry.initializers("zero_init.v0")
def configure_zero_init(*, inplace: bool = False) -> Callable[[Array], Array]:
    return partial(zero_init, inplace=inplace)


def uniform_init(
    data: Array, *, lo: float = -0.1, hi: float = 0.1, inplace: bool = False
) -> Array:
    xp = get_array_module(data)
    values = xp.random.uniform(lo, hi, data.shape)
    if inplace:
        copy_array(data, values)
        return data
    else:
        return values.astype(data.dtype)


@registry.initializers("uniform_init.v0")
def configure_uniform_init(
    *, lo: float = -0.1, hi: float = 0.1, inplace: bool = False
) -> Callable[[Array], Array]:
    return partial(uniform_init, lo=lo, hi=hi, inplace=inplace)


def normal_init(data: Array, *, fan_in: int = -1, inplace: bool = False) -> Array:
    if fan_in == -1:
        fan_in = data.shape[1]
    xp = get_array_module(data)
    scale = xp.sqrt(1.0 / fan_in)
    size = int(xp.prod(data.shape))
    inits = xp.random.normal(scale=scale, size=size)
    inits = inits.reshape(data.shape)
    if inplace:
        copy_array(data, inits)
        return data
    else:
        return inits


@registry.initializers("normal_init.v0")
def configure_normal_init(
    *, fan_in: int = -1, inplace: bool = False
) -> Callable[[Array], Array]:
    return partial(normal_init, fan_in=fan_in, inplace=inplace)


__all__ = ["normal_init", "uniform_init", "xavier_uniform_init", "zero_init"]<commit_msg>Adjust initializer functions to be Jax-compatible<commit_after>from typing import Callable

from .backends import Ops
from .config import registry
from .types import Array, Shape
from .util import copy_array, partial


def xavier_uniform_init(ops: Ops, shape: Shape) -> Array:
    scale = ops.xp.sqrt(6.0 / (shape[0] + shape[1]))
    return ops.xp.random.uniform(-scale, scale, shape).astype("float32")


@registry.initializers("xavier_uniform_init.v0")
def configure_xavier_uniform_init(*, inplace: bool = False) -> Callable[[Shape], Array]:
    return partial(xavier_uniform_init, inplace=inplace)


def zero_init(ops: Ops, shape: Shape) -> Array:
    return ops.alloc(*shape)


@registry.initializers("zero_init.v0")
def configure_zero_init(*, inplace: bool = False) -> Callable[[Array], Array]:
    return partial(zero_init, inplace=inplace)


def uniform_init(
        ops: Ops, shape: Shape, *, lo: float = -0.1, hi: float = 0.1
) -> Array:
    values = ops.xp.random.uniform(lo, hi, shape)
    return values.astype(data.dtype)


@registry.initializers("uniform_init.v0")
def configure_uniform_init(
    *, lo: float = -0.1, hi: float = 0.1, inplace: bool = False
) -> Callable[[Array], Array]:
    return partial(uniform_init, lo=lo, hi=hi, inplace=inplace)


def normal_init(ops: Ops, shape: Shape, *, fan_in: int = -1, inplace: bool = False) -> Array:
    if fan_in == -1:
        fan_in = shape[1]
    scale = ops.xp.sqrt(1.0 / fan_in)
    size = int(ops.xp.prod(data.shape))
    inits = ops.xp.random.normal(scale=scale, size=size).astype("float32")
    inits = inits.reshape(data.shape)
    return inits


@registry.initializers("normal_init.v0")
def configure_normal_init(
    *, fan_in: int = -1, inplace: bool = False
) -> Callable[[Array], Array]:
    return partial(normal_init, fan_in=fan_in, inplace=inplace)


__all__ = ["normal_init", "uniform_init", "xavier_uniform_init", "zero_init"]
{'fraction_non_alphanumeric': 0.1388888888888889, 'fraction_numerical': 0.011770244821092278, 'mean_word_length': 3.828409090909091, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7519512', 'n_tokens_mistral': 1499, 'n_tokens_neox': 1467, 'n_words': 344}