diff --git a/tensor-shapes/pyrefly-jax-stubs/jax-stubs/_shapes.pyi b/tensor-shapes/pyrefly-jax-stubs/jax-stubs/_shapes.pyi index bb2cccb4c6..7173a007e0 100644 --- a/tensor-shapes/pyrefly-jax-stubs/jax-stubs/_shapes.pyi +++ b/tensor-shapes/pyrefly-jax-stubs/jax-stubs/_shapes.pyi @@ -2224,3 +2224,40 @@ def triu_indices_from_shape(shape: IntTuple, k: int) -> IntTuple: s0 = shape[0] s1 = shape[1] return triu_indices_shape(s0, k, s1) + +@type_shape_dsl_function +def one_hot_shape(shape: IntTuple, num_classes: Int, axis: int) -> IntTuple: + zero_tuple = dsl.IntTuple((0,)) + zero = zero_tuple[0] + if dsl.is_concrete_int(num_classes) and num_classes < zero: + return dsl.Invalid("num_classes must be non-negative") + out_rank = len(shape) + 1 + if axis < 0 - out_rank or axis >= out_rank: + return dsl.Invalid("axis out of bounds") + if axis < 0: + norm_axis = axis + out_rank + else: + norm_axis = axis + 0 + return dsl.concat( + dsl.concat(shape[:norm_axis], dsl.IntTuple((num_classes,))), + shape[norm_axis:], + ) + +@type_shape_dsl_function +def glu_shape(shape: IntTuple, axis: int) -> IntTuple: + rank = len(shape) + if rank == 0: + return dsl.Invalid("glu requires at least 1-D array") + if axis < 0 - rank or axis >= rank: + return dsl.Invalid("axis out of bounds") + if axis < 0: + norm_axis = axis + rank + else: + norm_axis = axis + 0 + extent = shape[norm_axis] + if dsl.is_concrete_int(extent) and extent % 2 != 0: + return dsl.Invalid("glu input dimension must be divisible by 2") + halved = extent // 2 + return dsl.IntTuple( + (halved if index == norm_axis else shape[index] for index in range(rank)) + ) diff --git a/tensor-shapes/pyrefly-jax-stubs/jax-stubs/nn/__init__.pyi b/tensor-shapes/pyrefly-jax-stubs/jax-stubs/nn/__init__.pyi index 6fd101a5c6..f6545569c3 100644 --- a/tensor-shapes/pyrefly-jax-stubs/jax-stubs/nn/__init__.pyi +++ b/tensor-shapes/pyrefly-jax-stubs/jax-stubs/nn/__init__.pyi @@ -3,48 +3,267 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from typing import Any, overload, Sequence +from typing import Any, Literal, overload, Sequence from jax._array import Array as _Array, ArrayLike as _ArrayLike -from shape_extensions import broadcast, IntTuple +from jax._shapes import glu_shape, one_hot_shape, reduce_shape +from shape_extensions import broadcast, Flag, Int, IntTuple, IntVar + +from . import initializers as initializers type _Shape = IntTuple +type _Axis = int | tuple[int, ...] | None # Every activation JAX exposes here is elementwise, so the shape is preserved. -# `softmax` and `log_softmax` normalize along `axis` but do not reduce it away. -# `elu` and `leaky_relu` take their parameter as an array as well as a scalar, -# and an array parameter broadcasts against the input rather than preserving -# its shape. +# `softmax`, `log_softmax`, and `standardize` normalize along `axis` but do not reduce it away. +# `elu`, `leaky_relu`, `celu`, and `squareplus` take their parameter as an array as well as a scalar, +# and an array parameter broadcasts against the input rather than preserving its shape. +# `glu` splits the specified axis in half. +# `logsumexp` and `logmeanexp` perform reductions over the specified axis. +# `one_hot` inserts a new dimension with `num_classes` at the specified axis. def relu[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... def relu6[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... def sigmoid[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... def softplus[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def sparse_plus[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def sparse_sigmoid[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... def soft_sign[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... def silu[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... def swish[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... -@overload -def elu[Shape: _Shape = []]( - x: _ArrayLike[Shape], alpha: float = ... -) -> _Array[Shape]: ... -@overload +def log_sigmoid[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def hard_sigmoid[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def hard_silu[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def hard_swish[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def hard_tanh[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def tanh[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def selu[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def mish[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def identity[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +def log1mexp[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... def elu[Shape: _Shape = [], ParamShape: _Shape = []]( x: _ArrayLike[Shape], alpha: _ArrayLike[ParamShape] = ... ) -> _Array[broadcast(Shape, ParamShape)]: ... -@overload -def leaky_relu[Shape: _Shape = []]( - x: _ArrayLike[Shape], negative_slope: float = ... -) -> _Array[Shape]: ... -@overload def leaky_relu[Shape: _Shape = [], ParamShape: _Shape = []]( x: _ArrayLike[Shape], negative_slope: _ArrayLike[ParamShape] = ... ) -> _Array[broadcast(Shape, ParamShape)]: ... +def celu[Shape: _Shape = [], ParamShape: _Shape = []]( + x: _ArrayLike[Shape], alpha: _ArrayLike[ParamShape] = ... +) -> _Array[broadcast(Shape, ParamShape)]: ... +def squareplus[Shape: _Shape = [], ParamShape: _Shape = []]( + x: _ArrayLike[Shape], b: _ArrayLike[ParamShape] = ... +) -> _Array[broadcast(Shape, ParamShape)]: ... def gelu[Shape: _Shape = []]( x: _ArrayLike[Shape], approximate: bool = ... ) -> _Array[Shape]: ... -def hard_tanh[Shape: _Shape = []](x: _ArrayLike[Shape]) -> _Array[Shape]: ... +@overload +def glu[Shape: _Shape = [], Axis: Flag[int] = -1]( + x: _ArrayLike[Shape], axis: Axis = -1 +) -> _Array[glu_shape(Shape, Axis)]: ... +@overload +def glu[Shape: _Shape = []]( + x: _ArrayLike[Shape], axis: int = -1 +) -> _Array[IntTuple]: ... def softmax[Shape: _Shape = []]( x: _ArrayLike[Shape], axis: Sequence[int] | int | None = ..., where: Any = ... ) -> _Array[Shape]: ... def log_softmax[Shape: _Shape = []]( x: _ArrayLike[Shape], axis: Sequence[int] | int | None = ..., where: Any = ... ) -> _Array[Shape]: ... +def standardize[Shape: _Shape = []]( + x: _ArrayLike[Shape], + axis: tuple[int, ...] | int | None = -1, + mean: _ArrayLike[Any] | None = None, + variance: _ArrayLike[Any] | None = None, + epsilon: float = 1e-5, + where: Any = None, +) -> _Array[Shape]: ... +@overload +def logsumexp[ + Axis: Flag[_Axis], + KeepDims: Flag[bool], + Shape: _Shape = [], +]( + a: _ArrayLike[Shape], + axis: Axis = None, + b: _ArrayLike[Any] | None = None, + keepdims: KeepDims = False, + return_sign: Literal[False] = False, + where: Any = None, +) -> _Array[reduce_shape(Shape, Axis, KeepDims)]: ... +@overload +def logsumexp[ + Axis: Flag[_Axis], + KeepDims: Flag[bool], + Shape: _Shape = [], +]( + a: _ArrayLike[Shape], + axis: Axis = None, + b: _ArrayLike[Any] | None = None, + keepdims: KeepDims = False, + return_sign: Literal[True] = ..., + where: Any = None, +) -> tuple[ + _Array[reduce_shape(Shape, Axis, KeepDims)], + _Array[reduce_shape(Shape, Axis, KeepDims)], +]: ... +@overload +def logsumexp[ + Shape: _Shape = [], +]( + a: _ArrayLike[Shape], + axis: Sequence[int] = ..., + b: _ArrayLike[Any] | None = None, + keepdims: bool = False, + return_sign: Literal[False] = False, + where: Any = None, +) -> _Array[IntTuple]: ... +@overload +def logsumexp[ + Shape: _Shape = [], +]( + a: _ArrayLike[Shape], + axis: Sequence[int] = ..., + b: _ArrayLike[Any] | None = None, + keepdims: bool = False, + return_sign: Literal[True] = ..., + where: Any = None, +) -> tuple[_Array[IntTuple], _Array[IntTuple]]: ... +@overload +def logsumexp( + a: _ArrayLike[Any], + axis: Sequence[int] | int | None = None, + b: _ArrayLike[Any] | None = None, + keepdims: bool = False, + return_sign: bool = False, + where: Any = None, +) -> _Array[IntTuple] | tuple[_Array[IntTuple], _Array[IntTuple]]: ... +@overload +def logmeanexp[ + Axis: Flag[_Axis], + KeepDims: Flag[bool], + Shape: _Shape = [], +]( + x: _ArrayLike[Shape], + axis: Axis = None, + where: Any = None, + keepdims: KeepDims = False, +) -> _Array[reduce_shape(Shape, Axis, KeepDims)]: ... +@overload +def logmeanexp[ + Shape: _Shape = [], +]( + x: _ArrayLike[Shape], + axis: Sequence[int] = ..., + where: Any = None, + keepdims: bool = False, +) -> _Array[IntTuple]: ... +@overload +def logmeanexp( + x: _ArrayLike[Any], + axis: Sequence[int] | int | None = None, + where: Any = None, + keepdims: bool = False, +) -> _Array[IntTuple]: ... +@overload +def one_hot[ + Shape: _Shape = [], + NumClasses: IntVar = 0, + Axis: Flag[int] = -1, +]( + x: _ArrayLike[Shape], + num_classes: Int[NumClasses], + *, + dtype: Any = None, + axis: Axis = -1, + out_sharding: Any = None, +) -> _Array[one_hot_shape(Shape, Int[NumClasses], Axis)]: ... +@overload +def one_hot[ + Shape: _Shape = [], +]( + x: _ArrayLike[Shape], + num_classes: int, + *, + dtype: Any = None, + axis: int = -1, + out_sharding: Any = None, +) -> _Array[IntTuple]: ... +@overload +def dot_product_attention[ + QueryShape: _Shape = [], +]( + query: _ArrayLike[QueryShape], + key: _ArrayLike[Any], + value: _ArrayLike[Any], + bias: _ArrayLike[Any] | None = None, + mask: _ArrayLike[Any] | None = None, + scale: float | None = None, + is_causal: bool = False, + query_seq_lengths: _ArrayLike[Any] | None = None, + key_value_seq_lengths: _ArrayLike[Any] | None = None, + local_window_size: int | tuple[int, int] | None = None, + return_residual: Literal[False] = False, + implementation: str | None = None, +) -> _Array[QueryShape]: ... +@overload +def dot_product_attention[ + QueryShape: _Shape = [], +]( + query: _ArrayLike[QueryShape], + key: _ArrayLike[Any], + value: _ArrayLike[Any], + bias: _ArrayLike[Any] | None = None, + mask: _ArrayLike[Any] | None = None, + scale: float | None = None, + is_causal: bool = False, + query_seq_lengths: _ArrayLike[Any] | None = None, + key_value_seq_lengths: _ArrayLike[Any] | None = None, + local_window_size: int | tuple[int, int] | None = None, + return_residual: Literal[True] = ..., + implementation: str | None = None, +) -> tuple[_Array[QueryShape], _Array[IntTuple]]: ... +@overload +def dot_product_attention( + query: _ArrayLike[Any], + key: _ArrayLike[Any], + value: _ArrayLike[Any], + bias: _ArrayLike[Any] | None = None, + mask: _ArrayLike[Any] | None = None, + scale: float | None = None, + is_causal: bool = False, + query_seq_lengths: _ArrayLike[Any] | None = None, + key_value_seq_lengths: _ArrayLike[Any] | None = None, + local_window_size: int | tuple[int, int] | None = None, + return_residual: bool = False, + implementation: str | None = None, +) -> _Array[IntTuple] | tuple[_Array[IntTuple], _Array[IntTuple]]: ... +@overload +def scaled_matmul[ + Batch: IntTuple, + M: IntVar, + N: IntVar, + K: IntVar, +]( + lhs: _ArrayLike[[*Batch, M, K]], + rhs: _ArrayLike[[*Batch, N, K]], + lhs_scales: _ArrayLike[Any], + rhs_scales: _ArrayLike[Any], + preferred_element_type: Any = None, +) -> _Array[[*Batch, M, N]]: ... +@overload +def scaled_matmul( + lhs: _ArrayLike[Any], + rhs: _ArrayLike[Any], + lhs_scales: _ArrayLike[Any], + rhs_scales: _ArrayLike[Any], + preferred_element_type: Any = None, +) -> _Array[IntTuple]: ... +def scaled_dot_general( + lhs: _ArrayLike[Any], + rhs: _ArrayLike[Any], + dimension_numbers: Any, + lhs_scale: _ArrayLike[Any] | None = None, + rhs_scale: _ArrayLike[Any] | None = None, + **kwargs: Any, +) -> _Array[IntTuple]: ... +def get_scaled_dot_general_config(mode: str, **kwargs: Any) -> Any: ... diff --git a/tensor-shapes/pyrefly-jax-stubs/jax-stubs/nn/initializers.pyi b/tensor-shapes/pyrefly-jax-stubs/jax-stubs/nn/initializers.pyi new file mode 100644 index 0000000000..6f55dc082b --- /dev/null +++ b/tensor-shapes/pyrefly-jax-stubs/jax-stubs/nn/initializers.pyi @@ -0,0 +1,77 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Any, Callable + +from jax._array import Array as _Array +from shape_extensions import IntTuple + +type _Shape = IntTuple + +type Initializer = Callable[..., _Array[Any]] + +def constant[Shape: _Shape = []]( + value: Any, dtype: Any = ... +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def uniform[Shape: _Shape = []]( + scale: float = ..., dtype: Any = ... +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def normal[Shape: _Shape = []]( + stddev: float = ..., dtype: Any = ... +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def truncated_normal[Shape: _Shape = []]( + lower: float = ..., upper: float = ..., dtype: Any = ... +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def zeros[Shape: _Shape = []]( + key: Any, shape: Shape, dtype: Any = ..., out_sharding: Any = ... +) -> _Array[Shape]: ... +def ones[Shape: _Shape = []]( + key: Any, shape: Shape, dtype: Any = ..., out_sharding: Any = ... +) -> _Array[Shape]: ... +def glorot_normal[Shape: _Shape = []]( + **kwargs: Any, +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def glorot_uniform[Shape: _Shape = []]( + **kwargs: Any, +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def he_normal[Shape: _Shape = []]( + **kwargs: Any, +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def he_uniform[Shape: _Shape = []]( + **kwargs: Any, +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def kaiming_normal[Shape: _Shape = []]( + **kwargs: Any, +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def kaiming_uniform[Shape: _Shape = []]( + **kwargs: Any, +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def lecun_normal[Shape: _Shape = []]( + **kwargs: Any, +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def lecun_uniform[Shape: _Shape = []]( + **kwargs: Any, +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def orthogonal[Shape: _Shape = []]( + scale: float = ..., column_axis: int = ..., dtype: Any = ... +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def delta_orthogonal[Shape: _Shape = []]( + scale: float = ..., column_axis: int = ..., dtype: Any = ... +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def variance_scaling[Shape: _Shape = []]( + scale: float = ..., + mode: str = ..., + distribution: str = ..., + in_axis: int | tuple[int, ...] = ..., + out_axis: int | tuple[int, ...] = ..., + batch_axis: tuple[int, ...] = ..., + dtype: Any = ..., +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def xavier_normal[Shape: _Shape = []]( + **kwargs: Any, +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... +def xavier_uniform[Shape: _Shape = []]( + **kwargs: Any, +) -> Callable[[Any, Shape, Any], _Array[Shape]]: ... diff --git a/tensor-shapes/pyrefly-jax-stubs/test/test_nn.py b/tensor-shapes/pyrefly-jax-stubs/test/test_nn.py index 5cef1acb86..2792fe007d 100644 --- a/tensor-shapes/pyrefly-jax-stubs/test/test_nn.py +++ b/tensor-shapes/pyrefly-jax-stubs/test/test_nn.py @@ -8,7 +8,7 @@ import jax import jax.nn as jnn import jax.numpy as jnp -from shape_extensions import assert_shape, IntVar +from shape_extensions import assert_shape, Int, IntVar N = IntVar("N") M = IntVar("M") @@ -76,3 +76,128 @@ def test_activations_accept_scalars() -> None: assert_shape(jnn.gelu(1.0).shape, ()) assert_shape(jnn.elu(-1.0, 0.5).shape, ()) assert_shape(jnn.leaky_relu(-1.0, 0.2).shape, ()) + + +def generic_one_hot[N: IntVar, K: IntVar]( + x: jax.Array[[N]], k: Int[K] +) -> jax.Array[[N, K]]: + return jnn.one_hot(x, k) + + +def test_one_hot() -> None: + x_1d = jnp.full((3,), 1) + assert_shape(jnn.one_hot(x_1d, 5).shape, (3, 5)) + assert_shape(jnn.one_hot(x_1d, 5, axis=0).shape, (5, 3)) + assert_shape(jnn.one_hot(x_1d, 5, axis=-1).shape, (3, 5)) + + x_2d = jnp.full((3, 4), 2) + assert_shape(jnn.one_hot(x_2d, 10).shape, (3, 4, 10)) + assert_shape(jnn.one_hot(x_2d, 10, axis=0).shape, (10, 3, 4)) + assert_shape(jnn.one_hot(x_2d, 10, axis=1).shape, (3, 10, 4)) + assert_shape(jnn.one_hot(x_2d, 10, axis=-1).shape, (3, 4, 10)) + + assert_shape(jnn.one_hot(1, 5).shape, (5,)) + assert_shape(jnn.one_hot(1, 5, axis=0).shape, (5,)) + assert_shape(jnn.one_hot(1, 5, axis=-1).shape, (5,)) + + +def test_more_elementwise_activations_preserve_shape() -> None: + x = jnp.full((3, 4), 0.5) + + assert_shape(jnn.identity(x).shape, (3, 4)) + assert_shape(jnn.sparse_plus(x).shape, (3, 4)) + assert_shape(jnn.sparse_sigmoid(x).shape, (3, 4)) + assert_shape(jnn.log_sigmoid(x).shape, (3, 4)) + assert_shape(jnn.hard_sigmoid(x).shape, (3, 4)) + assert_shape(jnn.hard_silu(x).shape, (3, 4)) + assert_shape(jnn.hard_swish(x).shape, (3, 4)) + assert_shape(jnn.selu(x).shape, (3, 4)) + assert_shape(jnn.mish(x).shape, (3, 4)) + assert_shape(jnn.tanh(x).shape, (3, 4)) + assert_shape(jnn.log1mexp(x).shape, (3, 4)) + + assert_shape(jnn.identity(0.5).shape, ()) + assert_shape(jnn.tanh(0.5).shape, ()) + assert_shape(jnn.selu(0.5).shape, ()) + assert_shape(jnn.mish(0.5).shape, ()) + + +def test_more_parameterized_activations_preserve_shape() -> None: + x = jnp.full((2, 5), -0.5) + + assert_shape(jnn.celu(x).shape, (2, 5)) + assert_shape(jnn.celu(x, 0.5).shape, (2, 5)) + assert_shape(jnn.squareplus(x).shape, (2, 5)) + assert_shape(jnn.squareplus(x, 2.0).shape, (2, 5)) + + +def test_more_parameterized_activations_broadcast_parameters() -> None: + x = jnp.full((3, 4), -0.5) + parameter = jnp.full((5, 1, 4), 0.2) + + assert_shape(jnn.celu(x, parameter).shape, (5, 3, 4)) + assert_shape(jnn.squareplus(x, parameter).shape, (5, 3, 4)) + + +def test_glu_splits_axis() -> None: + x = jnp.ones((4, 6)) + + assert_shape(jnn.glu(x).shape, (4, 3)) + assert_shape(jnn.glu(x, -1).shape, (4, 3)) + assert_shape(jnn.glu(x, 0).shape, (2, 6)) + assert_shape(jnn.glu(x, axis=1).shape, (4, 3)) + + +def test_standardize_normalizes_without_reducing() -> None: + x = jnp.full((3, 4), 0.25) + + assert_shape(jnn.standardize(x).shape, (3, 4)) + assert_shape(jnn.standardize(x, 0).shape, (3, 4)) + assert_shape(jnn.standardize(x, -1).shape, (3, 4)) + assert_shape(jnn.standardize(x, (0,)).shape, (3, 4)) + + +def test_logsumexp_and_logmeanexp() -> None: + x = jnp.full((3, 4), 0.5) + + assert_shape(jnn.logsumexp(x).shape, ()) + assert_shape(jnn.logsumexp(x, 0).shape, (4,)) + assert_shape(jnn.logsumexp(x, -1).shape, (3,)) + assert_shape(jnn.logsumexp(x, 0, keepdims=True).shape, (1, 4)) + + res, sign = jnn.logsumexp(x, 0, return_sign=True) + assert_shape(res.shape, (4,)) + assert_shape(sign.shape, (4,)) + + assert_shape(jnn.logmeanexp(x).shape, ()) + assert_shape(jnn.logmeanexp(x, 0).shape, (4,)) + assert_shape(jnn.logmeanexp(x, -1).shape, (3,)) + assert_shape(jnn.logmeanexp(x, 0, keepdims=True).shape, (1, 4)) + + +def test_dot_product_attention() -> None: + q = jnp.ones((2, 8, 4, 16)) + k = jnp.ones((2, 8, 4, 16)) + v = jnp.ones((2, 8, 4, 16)) + + assert_shape(jnn.dot_product_attention(q, k, v).shape, (2, 8, 4, 16)) + out, res = jnn.dot_product_attention(q, k, v, return_residual=True) + assert_shape(out.shape, (2, 8, 4, 16)) + + +def test_scaled_matmul() -> None: + lhs = jnp.ones((2, 4, 8)) + rhs = jnp.ones((2, 6, 8)) + lhs_s = jnp.ones((2, 4, 1)) + rhs_s = jnp.ones((2, 6, 1)) + + assert_shape(jnn.scaled_matmul(lhs, rhs, lhs_s, rhs_s).shape, (2, 4, 6)) + + +def test_initializers() -> None: + key = jax.random.key(0) + + assert_shape(jnn.initializers.zeros(key, (3, 4)).shape, (3, 4)) + assert_shape(jnn.initializers.ones(key, (2, 5)).shape, (2, 5)) + init = jnn.initializers.normal() + assert_shape(init(key, (4, 4), jnp.float32).shape, (4, 4))