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"]

@@ -22,16 +22,42 @@
 llable%0A%0A
+from .backends import Ops%0A
 from .co
@@ -101,16 +101,23 @@
 rt Array
+, Shape
 %0Afrom .u
@@ -130,26 +130,8 @@
 port
- get_array_module,
  cop
@@ -177,88 +177,41 @@
 nit(
-data: Array, *, inplace: bool = False) -%3E Array:%0A    xp = get_array_module(data)
+ops: Ops, shape: Shape) -%3E Array:
 %0A   
@@ -211,32 +211,36 @@
 ay:%0A    scale = 
+ops.
 xp.sqrt(6.0 / (d
@@ -238,21 +238,16 @@
 (6.0 / (
-data.
 shape%5B0%5D
@@ -249,21 +249,16 @@
 pe%5B0%5D + 
-data.
 shape%5B1%5D
@@ -268,45 +268,19 @@
     
-if inplace:%0A        copy_array(data, 
+return ops.
 xp.r
@@ -312,109 +312,31 @@
 le, 
-data.
 shape)
-)%0A        return data%0A    else:%0A        return xp.random.uniform(-scale, scale, data.shape
+.astype(%22float32%22
 )%0A%0A%0A
@@ -449,37 +449,37 @@
 e) -%3E Callable%5B%5B
-Array
+Shape
 %5D, Array%5D:%0A    r
@@ -550,195 +550,69 @@
 nit(
-data: Array, *, inplace: bool = False) -%3E Array:%0A    if inplace:%0A        data.fill(0.0)%0A        return data%0A    else:%0A        xp = get_array_module(data)%0A        return xp.zeros_like(data
+ops: Ops, shape: Shape) -%3E Array:%0A    return ops.alloc(*shape
 )%0A%0A%0A
@@ -800,27 +800,42 @@
 it(%0A    
-data: Array
+    ops: Ops, shape: Shape
 , *, lo:
@@ -868,75 +868,20 @@
  0.1
-, inplace: bool = False%0A) -%3E Array:%0A    xp = get_array_module(data)
+%0A) -%3E Array:
 %0A   
@@ -886,24 +886,28 @@
    values = 
+ops.
 xp.random.un
@@ -924,103 +924,15 @@
 hi, 
-data.
 shape)%0A
-    if inplace:%0A        copy_array(data, values)%0A        return data%0A    else:%0A    
     
@@ -1217,19 +1217,30 @@
 nit(
-data: Array
+ops: Ops, shape: Shape
 , *,
@@ -1329,21 +1329,16 @@
 an_in = 
-data.
 shape%5B1%5D
@@ -1342,40 +1342,8 @@
 %5B1%5D%0A
-    xp = get_array_module(data)%0A
     
@@ -1350,16 +1350,20 @@
 scale = 
+ops.
 xp.sqrt(
@@ -1391,16 +1391,20 @@
 e = int(
+ops.
 xp.prod(
@@ -1428,16 +1428,20 @@
 inits = 
+ops.
 xp.rando
@@ -1472,16 +1472,34 @@
 ze=size)
+.astype(%22float32%22)
 %0A    ini
@@ -1533,90 +1533,8 @@
 pe)%0A
-    if inplace:%0A        copy_array(data, inits)%0A        return data%0A    else:%0A    
     

Adjust initializer functions to be Jax-compatible