From e542413e7338427b1776b177947995e6836390c9 Mon Sep 17 00:00:00 2001 From: ChenAo-Phys Date: Fri, 21 Jun 2024 10:24:59 +0200 Subject: [PATCH 1/5] add support for complex networks --- equinox/_misc.py | 19 +++++++++++++++- equinox/nn/_conv.py | 52 +++++++++++-------------------------------- equinox/nn/_linear.py | 20 +++++------------ equinox/nn/_rnn.py | 38 +++++++++++-------------------- 4 files changed, 50 insertions(+), 79 deletions(-) diff --git a/equinox/_misc.py b/equinox/_misc.py index 14183278..63f5f93b 100644 --- a/equinox/_misc.py +++ b/equinox/_misc.py @@ -1,7 +1,10 @@ +from typing import Any + import jax import jax.core import jax.numpy as jnp -from jaxtyping import Array +import jax.random as jrandom +from jaxtyping import Array, PRNGKeyArray def left_broadcast_to(arr: Array, shape: tuple[int, ...]) -> Array: @@ -18,3 +21,17 @@ def default_floating_dtype(): return jnp.float64 else: return jnp.float32 + + +def default_init( + key: PRNGKeyArray, shape: tuple[int, ...], dtype: Any, lim: float +) -> jax.Array: + if jnp.issubdtype(dtype, jnp.complexfloating): + # only two possible complex dtypes, jnp.complex64 or jnp.complex128 + real_dtype = jnp.float32 if dtype == jnp.complex64 else jnp.float64 + rkey, ikey = jrandom.split(key, 2) + real = jrandom.uniform(rkey, shape, real_dtype, minval=-lim, maxval=lim) + imag = jrandom.uniform(ikey, shape, real_dtype, minval=-lim, maxval=lim) + return real.astype(dtype) + 1j * imag.astype(dtype) + else: + return jrandom.uniform(key, shape, dtype, minval=-lim, maxval=lim) diff --git a/equinox/nn/_conv.py b/equinox/nn/_conv.py index db0bb351..aceea590 100644 --- a/equinox/nn/_conv.py +++ b/equinox/nn/_conv.py @@ -10,7 +10,7 @@ import numpy as np from jaxtyping import Array, PRNGKeyArray -from .._misc import default_floating_dtype +from .._misc import default_floating_dtype, default_init from .._module import field, Module from ._misc import all_sequences @@ -153,8 +153,6 @@ def __init__( case the padding is an odd number, then the extra padding is added at the end for `'SAME'` and at the beginning for `'SAME_LOWER'`. """ - dtype = default_floating_dtype() if dtype is None else dtype - wkey, bkey = jrandom.split(key, 2) parse = _ntuple(num_spatial_dims) kernel_size = parse(kernel_size) @@ -167,25 +165,14 @@ def __init__( f"by `groups` (={groups})." ) + dtype = default_floating_dtype() if dtype is None else dtype + wkey, bkey = jrandom.split(key, 2) grouped_in_channels = in_channels // groups - lim = 1 / np.sqrt(grouped_in_channels * math.prod(kernel_size)) - self.weight = jrandom.uniform( - wkey, - (out_channels, grouped_in_channels) + kernel_size, - minval=-lim, - maxval=lim, - dtype=dtype, - ) - if use_bias: - self.bias = jrandom.uniform( - bkey, - (out_channels,) + (1,) * num_spatial_dims, - minval=-lim, - maxval=lim, - dtype=dtype, - ) - else: - self.bias = None + lim = 1 / math.sqrt(grouped_in_channels * math.prod(kernel_size)) + wshape = (out_channels, grouped_in_channels) + kernel_size + self.weight = default_init(wkey, wshape, dtype, lim) + bshape = (out_channels,) + (1,) * num_spatial_dims + self.bias = default_init(bkey, bshape, dtype, lim) if use_bias else None self.num_spatial_dims = num_spatial_dims self.in_channels = in_channels @@ -520,24 +507,11 @@ def __init__( raise ValueError("Must have `output_padding < stride` (elementwise).") grouped_in_channels = in_channels // groups - lim = 1 / np.sqrt(grouped_in_channels * math.prod(kernel_size)) - self.weight = jrandom.uniform( - wkey, - (out_channels, grouped_in_channels) + kernel_size, - minval=-lim, - maxval=lim, - dtype=dtype, - ) - if use_bias: - self.bias = jrandom.uniform( - bkey, - (out_channels,) + (1,) * num_spatial_dims, - minval=-lim, - maxval=lim, - dtype=dtype, - ) - else: - self.bias = None + lim = 1 / math.sqrt(grouped_in_channels * math.prod(kernel_size)) + wshape = (out_channels, grouped_in_channels) + kernel_size + self.weight = default_init(wkey, wshape, dtype, lim) + bshape = (out_channels,) + (1,) * num_spatial_dims + self.bias = default_init(bkey, bshape, dtype, lim) if use_bias else None padding = _padding_init(padding, num_spatial_dims) padding_mode = _padding_mode_init(padding_mode) diff --git a/equinox/nn/_linear.py b/equinox/nn/_linear.py index 450ee663..c9654699 100644 --- a/equinox/nn/_linear.py +++ b/equinox/nn/_linear.py @@ -6,7 +6,7 @@ import jax.random as jrandom from jaxtyping import Array, PRNGKeyArray -from .._misc import default_floating_dtype +from .._misc import default_floating_dtype, default_init from .._module import field, Module @@ -47,23 +47,15 @@ def __init__( Likewise `out_features` can also be a string `"scalar"`, in which case the output from the layer will have shape `()`. """ + dtype = default_floating_dtype() if dtype is None else dtype wkey, bkey = jrandom.split(key, 2) in_features_ = 1 if in_features == "scalar" else in_features out_features_ = 1 if out_features == "scalar" else out_features lim = 1 / math.sqrt(in_features_) - - if dtype is None: - dtype = default_floating_dtype() - - self.weight = jrandom.uniform( - wkey, (out_features_, in_features_), minval=-lim, maxval=lim, dtype=dtype - ) - if use_bias: - self.bias = jrandom.uniform( - bkey, (out_features_,), minval=-lim, maxval=lim, dtype=dtype - ) - else: - self.bias = None + wshape = (out_features_, in_features_) + self.weight = default_init(wkey, wshape, dtype, lim) + bshape = (out_features_,) + self.bias = default_init(bkey, bshape, dtype, lim) if use_bias else None self.in_features = in_features self.out_features = out_features diff --git a/equinox/nn/_rnn.py b/equinox/nn/_rnn.py index a6332203..b7c87a5c 100644 --- a/equinox/nn/_rnn.py +++ b/equinox/nn/_rnn.py @@ -7,7 +7,7 @@ import jax.random as jrandom from jaxtyping import Array, PRNGKeyArray -from .._misc import default_floating_dtype +from .._misc import default_floating_dtype, default_init from .._module import field, Module @@ -66,19 +66,13 @@ def __init__( ihkey, hhkey, bkey, bkey2 = jrandom.split(key, 4) lim = math.sqrt(1 / hidden_size) - self.weight_ih = jrandom.uniform( - ihkey, (3 * hidden_size, input_size), minval=-lim, maxval=lim, dtype=dtype - ) - self.weight_hh = jrandom.uniform( - hhkey, (3 * hidden_size, hidden_size), minval=-lim, maxval=lim, dtype=dtype - ) + ihshape = (3 * hidden_size, input_size) + self.weight_ih = default_init(ihkey, ihshape, dtype, lim) + hhshape = (3 * hidden_size, hidden_size) + self.weight_hh = default_init(hhkey, hhshape, dtype, lim) if use_bias: - self.bias = jrandom.uniform( - bkey, (3 * hidden_size,), minval=-lim, maxval=lim, dtype=dtype - ) - self.bias_n = jrandom.uniform( - bkey2, (hidden_size,), minval=-lim, maxval=lim, dtype=dtype - ) + self.bias = default_init(bkey, (3 * hidden_size,), dtype, lim) + self.bias_n = default_init(bkey2, (hidden_size,), dtype, lim) else: self.bias = None self.bias_n = None @@ -172,18 +166,12 @@ def __init__( ihkey, hhkey, bkey = jrandom.split(key, 3) lim = math.sqrt(1 / hidden_size) - self.weight_ih = jrandom.uniform( - ihkey, (4 * hidden_size, input_size), minval=-lim, maxval=lim, dtype=dtype - ) - self.weight_hh = jrandom.uniform( - hhkey, (4 * hidden_size, hidden_size), minval=-lim, maxval=lim, dtype=dtype - ) - if use_bias: - self.bias = jrandom.uniform( - bkey, (4 * hidden_size,), minval=-lim, maxval=lim, dtype=dtype - ) - else: - self.bias = None + ihshape = (4 * hidden_size, input_size) + self.weight_ih = default_init(ihkey, ihshape, dtype, lim) + hhshape = (4 * hidden_size, hidden_size) + self.weight_hh = default_init(hhkey, hhshape, dtype, lim) + bshape = (4 * hidden_size,) + self.bias = default_init(bkey, bshape, dtype, lim) if use_bias else None self.input_size = input_size self.hidden_size = hidden_size From 9774d23e657b23d0b04f7810b7023f6b5beb5cbf Mon Sep 17 00:00:00 2001 From: ChenAo-Phys Date: Fri, 21 Jun 2024 10:25:51 +0200 Subject: [PATCH 2/5] add tests for complex networks --- tests/test_nn.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_nn.py b/tests/test_nn.py index 769bf145..157b799c 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -48,6 +48,10 @@ def test_linear(getkey): x = jrandom.normal(getkey(), (2,), dtype=jnp.float16) assert linear(x).dtype == jnp.float16 + linear = eqx.nn.Linear(2, "scalar", key=getkey(), dtype=jnp.complex64) + x = jrandom.normal(getkey(), (2,), dtype=jnp.complex64) + assert linear(x).dtype == jnp.complex64 + def test_identity(getkey): identity1 = eqx.nn.Identity() @@ -344,6 +348,18 @@ def test_conv2d(getkey): answer = jnp.array([-37, -31, -9, 25, 61, 49, 23, 41, 27]).reshape(1, 3, 3) assert jnp.allclose(conv(data), answer) + # Test complex value matches + conv = eqx.nn.Conv2d(1, 1, 3, padding=1, dtype=jnp.complex64, key=getkey()) + new_weight = jnp.arange(9, dtype=jnp.complex64).reshape(1, 1, 3, 3) + new_bias = jnp.array([1 + 1j], dtype=jnp.complex64).reshape(1, 1, 1) + data = (1 + 1j) * jnp.arange(-4, 5, dtype=jnp.complex64).reshape(1, 3, 3) + assert new_weight.shape == conv.weight.shape + assert new_bias.shape == conv.bias.shape # pyright: ignore + conv = eqx.tree_at(lambda x: (x.weight, x.bias), conv, (new_weight, new_bias)) + answer = jnp.array([-37, -31, -9, 25, 61, 49, 23, 41, 27]).reshape(1, 3, 3) + answer = (1 + 1j) * answer.astype(jnp.complex64) + assert jnp.allclose(conv(data), answer) + # Test groups conv = eqx.nn.Conv2d(2, 2, kernel_size=3, padding=1, key=getkey(), groups=2) # we will duplicate the weights from the "value matches" case From 05d20fa4d3feee6cd6189265881c1b9069296029 Mon Sep 17 00:00:00 2001 From: ChenAo-Phys Date: Fri, 21 Jun 2024 10:26:55 +0200 Subject: [PATCH 3/5] support for older version of jax --- tests/test_debug.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/test_debug.py b/tests/test_debug.py index 0778a5b6..bfa02b53 100644 --- a/tests/test_debug.py +++ b/tests/test_debug.py @@ -30,10 +30,9 @@ def f(x, terminate): f(jnp.array(1.0), terminate=False) jax.effects_barrier() text, _ = capfd.readouterr() - assert ( - text - == "foo:\n primals=Array(1., dtype=float32)\ncotangents=Array(nan, dtype=float32)\n" # noqa: E501 - ) + out_text1 = "foo:\n primals=Array(1., dtype=float32)\ncotangents=Array(nan, dtype=float32)\n" # noqa: E501 + out_text2 = "foo:\n primals=array(1., dtype=float32)\ncotangents=array(nan, dtype=float32)\n" # noqa: E501 + assert text in (out_text1, out_text2) with pytest.raises(Exception): f(jnp.array(1.0), terminate=True) From 08e114faf3f60e4ca91b1274c15438867c4a5daf Mon Sep 17 00:00:00 2001 From: ChenAo-Phys Date: Fri, 21 Jun 2024 14:22:31 +0200 Subject: [PATCH 4/5] replace PositionalSharding by NamedSharding --- tests/test_sharding.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_sharding.py b/tests/test_sharding.py index c22fdf1b..9d74f45b 100644 --- a/tests/test_sharding.py +++ b/tests/test_sharding.py @@ -2,10 +2,11 @@ import jax import jax.random as jr import jax.tree_util as jtu +from jax.sharding import Mesh, PartitionSpec [cpu] = jax.local_devices(backend="cpu") -sharding = jax.sharding.PositionalSharding([cpu]) +sharding = jax.sharding.NamedSharding(Mesh([cpu], "x"), PartitionSpec("x")) def test_sharding(): From 05538bb609c154b23003f9b5b997c164d8254d78 Mon Sep 17 00:00:00 2001 From: ChenAo-Phys Date: Sat, 22 Jun 2024 10:58:29 +0200 Subject: [PATCH 5/5] move default_init to nn._misc --- equinox/_misc.py | 19 +------------------ equinox/nn/_conv.py | 4 ++-- equinox/nn/_linear.py | 3 ++- equinox/nn/_misc.py | 18 ++++++++++++++++++ equinox/nn/_rnn.py | 3 ++- 5 files changed, 25 insertions(+), 22 deletions(-) diff --git a/equinox/_misc.py b/equinox/_misc.py index 63f5f93b..14183278 100644 --- a/equinox/_misc.py +++ b/equinox/_misc.py @@ -1,10 +1,7 @@ -from typing import Any - import jax import jax.core import jax.numpy as jnp -import jax.random as jrandom -from jaxtyping import Array, PRNGKeyArray +from jaxtyping import Array def left_broadcast_to(arr: Array, shape: tuple[int, ...]) -> Array: @@ -21,17 +18,3 @@ def default_floating_dtype(): return jnp.float64 else: return jnp.float32 - - -def default_init( - key: PRNGKeyArray, shape: tuple[int, ...], dtype: Any, lim: float -) -> jax.Array: - if jnp.issubdtype(dtype, jnp.complexfloating): - # only two possible complex dtypes, jnp.complex64 or jnp.complex128 - real_dtype = jnp.float32 if dtype == jnp.complex64 else jnp.float64 - rkey, ikey = jrandom.split(key, 2) - real = jrandom.uniform(rkey, shape, real_dtype, minval=-lim, maxval=lim) - imag = jrandom.uniform(ikey, shape, real_dtype, minval=-lim, maxval=lim) - return real.astype(dtype) + 1j * imag.astype(dtype) - else: - return jrandom.uniform(key, shape, dtype, minval=-lim, maxval=lim) diff --git a/equinox/nn/_conv.py b/equinox/nn/_conv.py index aceea590..9f634786 100644 --- a/equinox/nn/_conv.py +++ b/equinox/nn/_conv.py @@ -10,9 +10,9 @@ import numpy as np from jaxtyping import Array, PRNGKeyArray -from .._misc import default_floating_dtype, default_init +from .._misc import default_floating_dtype from .._module import field, Module -from ._misc import all_sequences +from ._misc import all_sequences, default_init _T = TypeVar("_T") diff --git a/equinox/nn/_linear.py b/equinox/nn/_linear.py index c9654699..4c8655de 100644 --- a/equinox/nn/_linear.py +++ b/equinox/nn/_linear.py @@ -6,8 +6,9 @@ import jax.random as jrandom from jaxtyping import Array, PRNGKeyArray -from .._misc import default_floating_dtype, default_init +from .._misc import default_floating_dtype from .._module import field, Module +from ._misc import default_init class Linear(Module, strict=True): diff --git a/equinox/nn/_misc.py b/equinox/nn/_misc.py index b5649970..63f0aba4 100644 --- a/equinox/nn/_misc.py +++ b/equinox/nn/_misc.py @@ -2,6 +2,11 @@ from collections.abc import Sequence from typing import Any, TYPE_CHECKING, TypeVar, Union +import jax +import jax.numpy as jnp +import jax.random as jrandom +from jaxtyping import PRNGKeyArray + _T = TypeVar("_T", bound=Sequence) @@ -17,3 +22,16 @@ def all_sequences( # beartype doesn't like StrictTypeGuard def all_sequences(x: Union[Sequence[Any], Sequence[_T]]) -> bool: return all(isinstance(xi, Sequence) for xi in x) + + +def default_init( + key: PRNGKeyArray, shape: tuple[int, ...], dtype: Any, lim: float +) -> jax.Array: + if jnp.issubdtype(dtype, jnp.complexfloating): + real_dtype = jnp.finfo(dtype).dtype + rkey, ikey = jrandom.split(key, 2) + real = jrandom.uniform(rkey, shape, real_dtype, minval=-lim, maxval=lim) + imag = jrandom.uniform(ikey, shape, real_dtype, minval=-lim, maxval=lim) + return real.astype(dtype) + 1j * imag.astype(dtype) + else: + return jrandom.uniform(key, shape, dtype, minval=-lim, maxval=lim) diff --git a/equinox/nn/_rnn.py b/equinox/nn/_rnn.py index b7c87a5c..e40f5c8f 100644 --- a/equinox/nn/_rnn.py +++ b/equinox/nn/_rnn.py @@ -7,8 +7,9 @@ import jax.random as jrandom from jaxtyping import Array, PRNGKeyArray -from .._misc import default_floating_dtype, default_init +from .._misc import default_floating_dtype from .._module import field, Module +from ._misc import default_init class GRUCell(Module, strict=True):