from typing import Callable
import numpy

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

# TODO: Harmonize naming with Keras, and fill in missing entries
# https://keras.io/initializers/ We should also have He normal/uniform
# and probably lecun normal/uniform.

# Initialize via numpy, before copying to ops. This makes it easier to work with
# the different backends, because the backend won't affect the randomization.


def lecun_normal_init(ops: Ops, shape: Shape) -> FloatsXd:
    scale = numpy.sqrt(1.0 / shape[1])
    return ops.asarray_f(numpy.random.normal(0, scale, shape))


@registry.initializers("lecun_normal_init.v1")
def configure_lecun_normal_init() -> Callable[[Shape], FloatsXd]:
    return partial(lecun_normal_init)


def he_normal_init(ops: Ops, shape: Shape) -> FloatsXd:
    scale = numpy.sqrt(2.0 / shape[1])
    return ops.asarray_f(numpy.random.normal(0, scale, shape))


@registry.initializers("he_normal_init.v1")
def configure_he_normal_init() -> Callable[[Shape], FloatsXd]:
    return partial(he_normal_init)


def glorot_normal_init(ops: Ops, shape: Shape) -> FloatsXd:
    scale = numpy.sqrt(2.0 / (shape[1] + shape[0]))
    return ops.asarray_f(numpy.random.normal(0, scale, shape))


@registry.initializers("glorot_normal_init.v1")
def configure_glorot_normal_init() -> Callable[[Shape], FloatsXd]:
    return partial(glorot_normal_init)


def he_uniform_init(ops: Ops, shape: Shape) -> FloatsXd:
    scale = numpy.sqrt(6.0 / shape[1])
    return ops.asarray_f(numpy.random.uniform(-scale, scale, shape))


@registry.initializers("he_uniform_init.v1")
def configure_he_uniform_init() -> Callable[[Shape], FloatsXd]:
    return partial(he_uniform_init)


def lecun_uniform_init(ops: Ops, shape: Shape) -> FloatsXd:
    scale = numpy.sqrt(3.0 / shape[1])
    return ops.asarray_f(numpy.random.uniform(-scale, scale, shape))


@registry.initializers("lecun_uniform_init.v1")
def configure_lecun_uniform_init() -> Callable[[Shape], FloatsXd]:
    return partial(lecun_uniform_init)


def glorot_uniform_init(ops: Ops, shape: Shape) -> FloatsXd:
    scale = numpy.sqrt(6.0 / (shape[0] + shape[1]))
    return ops.asarray_f(numpy.random.uniform(-scale, scale, shape))


@registry.initializers("glorot_uniform_init.v1")
def configure_glorot_uniform_init() -> Callable[[Shape], FloatsXd]:
    return partial(glorot_uniform_init)


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


@registry.initializers("zero_init.v1")
def configure_zero_init() -> Callable[[FloatsXd], FloatsXd]:
    return partial(zero_init)


def uniform_init(
    ops: Ops, shape: Shape, *, lo: float = -0.1, hi: float = 0.1
) -> FloatsXd:
    values = numpy.random.uniform(lo, hi, shape)
    return ops.asarray_f(values.astype("float32"))


@registry.initializers("uniform_init.v1")
def configure_uniform_init(
    *, lo: float = -0.1, hi: float = 0.1
) -> Callable[[FloatsXd], FloatsXd]:
    return partial(uniform_init, lo=lo, hi=hi)


def normal_init(ops: Ops, shape: Shape, *, mean: int = 0) -> FloatsXd:
    size = int(ops.xp.prod(ops.xp.asarray(shape)))
    inits = numpy.random.normal(scale=mean, size=size).astype("float32")
    inits = ops.reshape_f(inits, shape)
    return ops.asarray_f(inits)


@registry.initializers("normal_init.v1")
def configure_normal_init(*, mean: float = 0) -> Callable[[FloatsXd], FloatsXd]:
    return partial(normal_init, mean=mean)


__all__ = [
    "normal_init",
    "uniform_init",
    "glorot_uniform_init",
    "zero_init",
    "lecun_uniform_init",
    "he_uniform_init",
    "glorot_normal_init",
    "he_normal_init",
    "lecun_normal_init",
]

@@ -20,17 +20,23 @@
 Callable
+, cast
 %0A
-
 import n
@@ -608,32 +608,47 @@
 n ops.asarray_f(
+cast(FloatsXd, 
 numpy.random.nor
@@ -660,32 +660,33 @@
 , scale, shape))
+)
 %0A%0A%0A@registry.ini
@@ -937,32 +937,47 @@
 n ops.asarray_f(
+cast(FloatsXd, 
 numpy.random.nor
@@ -989,32 +989,33 @@
 , scale, shape))
+)
 %0A%0A%0A@registry.ini
@@ -1274,32 +1274,47 @@
 n ops.asarray_f(
+cast(FloatsXd, 
 numpy.random.nor
@@ -1326,32 +1326,33 @@
 , scale, shape))
+)
 %0A%0A%0A@registry.ini
@@ -1607,32 +1607,47 @@
 n ops.asarray_f(
+cast(FloatsXd, 
 numpy.random.uni
@@ -1665,32 +1665,33 @@
 , scale, shape))
+)
 %0A%0A%0A@registry.ini
@@ -1940,32 +1940,47 @@
 n ops.asarray_f(
+cast(FloatsXd, 
 numpy.random.uni
@@ -1998,32 +1998,33 @@
 , scale, shape))
+)
 %0A%0A%0A@registry.ini
@@ -2304,16 +2304,31 @@
 array_f(
+cast(FloatsXd, 
 numpy.ra
@@ -2362,16 +2362,17 @@
  shape))
+)
 %0A%0A%0A@regi
@@ -2910,16 +2910,31 @@
 array_f(
+cast(FloatsXd, 
 values.a
@@ -2950,16 +2950,17 @@
 oat32%22))
+)
 %0A%0A%0A@regi
@@ -3284,16 +3284,31 @@
  inits =
+ cast(FloatsXd,
  numpy.r
@@ -3360,16 +3360,17 @@
 loat32%22)
+)
 %0A    ini

Update initializers for typing in numpy 1.21+ (#516)