upload models

PirateNet.py

@@ -0,0 +1,85 @@
+import jax
+import jax.numpy as jnp
+import flax.linen as nn
+from .utils import Dense, FourierEmbs
+from typing import Union, Dict, Callable
+
+class PIModifiedBottleneck(nn.Module):
+    hidden_dim: int
+    output_dim: int
+    act: Callable
+    nonlinearity: float
+    reparam: Union[None, Dict]
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(self, x, u, v):
+        identity = x
+
+        x = Dense(features=self.hidden_dim, reparam=self.reparam, dtype=self.dtype)(x)
+        x = self.act(x)
+
+        x = x * u + (1 - x) * v
+
+        x = Dense(features=self.hidden_dim, reparam=self.reparam, dtype=self.dtype)(x)
+        x = self.act(x)
+
+        x = x * u + (1 - x) * v
+
+        x = Dense(features=self.output_dim, reparam=self.reparam, dtype=self.dtype)(x)
+        x = self.act(x)
+
+        alpha = self.param("alpha", nn.initializers.constant(self.nonlinearity), (1,))
+        x = alpha * x + (1 - alpha) * identity
+
+        return x
+
+class PirateNet(nn.Module):
+    num_layers: int
+    hidden_dim: int
+    output_dim: int
+    act: Callable = nn.silu
+    nonlinearity: float = 0.0
+    pi_init: Union[None, jnp.ndarray] = None
+    reparam : Union[None, Dict] = None
+    fourier_emb : Union[None, Dict] = None
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(self, x):
+        embs = FourierEmbs(**self.fourier_emb)(x)
+        x = embs
+
+        u = Dense(features=self.hidden_dim, reparam=self.reparam, dtype=self.dtype)(x)
+        u = self.act(u)
+
+        v = Dense(features=self.hidden_dim, reparam=self.reparam, dtype=self.dtype)(x)
+        v = self.act(v)
+
+        for _ in range(self.num_layers):
+            x = PIModifiedBottleneck(
+                hidden_dim=self.hidden_dim,
+                output_dim=x.shape[-1],
+                act=self.act,
+                nonlinearity=self.nonlinearity,
+                reparam=self.reparam,
+                dtype=self.dtype
+            )(x, u, v)
+
+        if self.pi_init is not None:
+            kernel = self.param("pi_init", nn.initializers.constant(self.pi_init, dtype=self.dtype), self.pi_init.shape)
+            y = jnp.dot(x, kernel)
+
+        else:
+            y = Dense(features=self.output_dim, reparam=self.reparam, dtype=self.dtype)(x)
+
+        return x, y
+    
+if __name__ == "__main__":
+    # Example usage
+    from activations import cauchy
+    cauchy_mod = lambda x : cauchy()(x)
+    model = PirateNet(num_layers=3, hidden_dim=32, output_dim=16, act=cauchy_mod, reparam=None, fourier_emb={'embed_scale': 1.0, 'embed_dim': 64})
+    params = model.init(jax.random.PRNGKey(0), jnp.ones(3))
+    output = model.apply(params, jnp.ones(3))
+    print(params)

__init__.py

@@ -0,0 +1,91 @@
+from .mlp_pinn import MLP_PINN
+from .PirateNet import PirateNet
+from .mlp import MLP
+from .siren import SIREN
+from .wire import WIRE
+from .activations import get_activation, list_activations
+
+model_key_dict = {
+    "MLP": MLP,
+    "SIREN": SIREN,
+    "WIRE": WIRE,
+    "PirateNet": PirateNet,
+    "MLP_PINN": MLP_PINN
+}
+
+def get_model(model_name : str):
+    """
+    Get the model class by name.
+
+    Args:
+        model_name (str): Name of the model.
+
+    Returns:
+        nn.Module: The model class.
+    """
+    if model_name not in model_key_dict:
+        raise ValueError(f"Model `{model_name}` is not supported. Supported models are: {list(model_key_dict.keys())}")
+    
+    return model_key_dict[model_name]
+
+def create_model_configs():
+    """
+    Create a dictionary of model configurations.
+
+    Returns:
+        dict: A dictionary of model configurations.
+    """
+    model_configs = {
+        "MLP": {},
+        "SIREN": {
+            "omega_0": 3.
+        },
+        "WIRE": {
+            "first_omega_0": 4.,
+            "hidden_omega_0": 4.,
+            "scale": 5.,
+        },
+        "PirateNet": {
+            "nonlinearity": 0.0,
+            "pi_init": None,
+            "reparam": {
+                "type": "weight_fact",
+                "mean": 1.0,
+                "stddev": 0.1,
+            },
+            "fourier_emb": {
+                "embed_scale": 2.,
+                "embed_dim": 256,
+            },
+        },
+        "MLP_PINN": {
+            "reparam": {
+                "type": "weight_fact",
+                "mean": 1.0,
+                "stddev": 0.1,
+            },
+            "fourier_emb": {
+                "embed_scale": 2.,
+                "embed_dim": 256,
+            },
+        },
+
+    }
+    return model_configs
+
+model_configs = create_model_configs()
+
+def get_extra_model_cfg(model_name: str):
+    """
+    Get the extra model configuration for a given model name.
+
+    Args:
+        model_name (str): Name of the model.
+
+    Returns:
+        dict: The extra model configuration.
+    """
+    if model_name not in model_configs:
+        raise ValueError(f"Model `{model_name}` is not supported. Available models are: {list(model_configs.keys())}")
+
+    return model_configs[model_name]

activations.py

@@ -0,0 +1,95 @@
+import jax
+import jax.numpy as jnp
+import flax.linen as nn
+
+activation_function = {
+    "relu": nn.relu,
+    "gelu": nn.gelu,
+    "silu": nn.silu,
+    "swish": nn.silu,
+    "tanh": nn.tanh,
+    "sigmoid": nn.sigmoid,
+    "softplus": nn.softplus,
+    "softmax": nn.softmax,
+    "leaky_relu": nn.leaky_relu,
+    "elu": nn.elu,
+    "selu": nn.selu,
+    "telu": lambda x: x * jnp.tanh(jnp.exp(x)),
+    "mish": lambda x: x * jnp.tanh(nn.softplus(x)),
+    "cauchy": lambda x: cauchy()(x),
+    "identity": lambda x: x,
+    "react": lambda x: react()(x),
+}
+
+# https://arxiv.org/pdf/2503.02267v1
+class react(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        a = self.param(
+            'a',
+            jax.nn.initializers.normal(0.1),
+            ()
+        )
+        b = self.param(
+            'b',
+            jax.nn.initializers.normal(0.1),
+            ()
+        )
+        
+        c = self.param(
+            'c',
+            jax.nn.initializers.normal(0.1),
+            ()
+        )
+        d = self.param(
+            'd',
+            jax.nn.initializers.normal(0.1),
+            ()
+        )
+
+        return (1 - jnp.exp(a * x + b)) / (1 + jnp.exp(c * x + d))
+
+# https://arxiv.org/abs/2409.19221
+class cauchy(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        l1 = self.param(
+            'lambda1',
+            jax.nn.initializers.constant(1.0),
+            ()
+        )
+        l2 = self.param(
+            'lambda2',
+            jax.nn.initializers.constant(1.0),
+            ()
+        )
+        d = self.param(
+            'd',
+            jax.nn.initializers.constant(1.0),
+            ()
+        )
+
+        return l1 * x / (x**2 + d**2) + l2 / (x**2 + d**2)
+
+def get_activation(name):
+    """
+    Get the activation function by name.
+
+    Args:
+        name (str): Name of the activation function.
+
+    Returns:
+        Callable: The activation function.
+    """
+    if name not in activation_function:
+        raise ValueError(f"Activation function `{name}` is not supported. Supported activations are : {list(activation_function.keys())}")
+    return activation_function[name]
+
+def list_activations():
+    """
+    List all available activation functions.
+
+    Returns:
+        list: A list of activation names.
+    """
+    return list(activation_function.keys())

mlp.py

@@ -0,0 +1,45 @@
+# Flax
+import jax.numpy as jnp
+from flax import linen as nn
+import jax
+from typing import Callable
+
+class MLP(nn.Module):
+    hidden_dim: int
+    output_dim: int
+    num_layers: int
+    act: Callable = nn.silu
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(self, x):
+        x = nn.Dense(
+            features=self.hidden_dim,
+            use_bias=True,
+            kernel_init=nn.initializers.glorot_normal(dtype=self.dtype),
+            param_dtype=self.dtype
+        )(x)
+        x = self.act(x)
+        for _ in range(self.num_layers):
+            x = nn.Dense(
+                features=self.hidden_dim,
+                use_bias=True,
+                kernel_init=nn.initializers.glorot_normal(dtype=self.dtype),
+                param_dtype=self.dtype
+            )(x)
+            x = self.act(x)
+        x = nn.Dense(
+            features=self.output_dim,
+            use_bias=True,
+            kernel_init=nn.initializers.glorot_normal(dtype=self.dtype),
+            param_dtype=self.dtype
+        )(x)
+        return x
+
+if __name__ == "__main__":
+    # Example usage
+    x = jax.random.uniform(jax.random.PRNGKey(0), (1, 3), minval=-3, maxval=3)
+    model = MLP(hidden_dim=32, output_dim=16, num_layers=3)
+    params = model.init(jax.random.PRNGKey(0), x)
+    model_fn = lambda params, x : model.apply(params, x)
+    print(model_fn(params, x).shape)

mlp_pinn.py

@@ -0,0 +1,44 @@
+# Flax
+import jax.numpy as jnp
+from flax import linen as nn
+from typing import Callable, Union, Dict
+from .utils import Dense, FourierEmbs
+
+# Modified MLP version based on the state-of-the-art practicies in PINN training:
+# Fourier embeddings and random weight factorization
+# You can read more about it in the paper: https://arxiv.org/pdf/2210.01274
+
+
+class MLP_PINN(nn.Module):
+    hidden_dim: int
+    output_dim: int
+    num_layers: int
+    act: Callable = nn.silu
+    dtype: jnp.dtype = jnp.float32
+    reparam : Union[None, Dict] = None
+    fourier_emb : Union[None, Dict] = None
+
+    @nn.compact
+    def __call__(self, x):
+        if self.fourier_emb is not None:
+            x = FourierEmbs(**self.fourier_emb)(x)
+        else:
+            x = Dense(
+                features=self.hidden_dim,
+                reparam=self.reparam,
+                dtype=self.dtype
+            )(x)
+            x = self.act(x)
+        for _ in range(self.num_layers):
+            x = Dense(
+                features=self.hidden_dim,
+                reparam=self.reparam,
+                dtype=self.dtype
+            )(x)
+            x = self.act(x)
+        x = Dense(
+            features=self.output_dim,
+            reparam=self.reparam,
+            dtype=self.dtype
+        )(x)
+        return x

siren.py

@@ -0,0 +1,66 @@
+import jax.numpy as jnp
+from flax import linen as nn
+
+from .utils import custom_uniform
+
+class SIREN(nn.Module):
+    output_dim: int
+    hidden_dim: int
+    num_layers: int
+    omega_0: float
+    dtype: jnp.dtype = jnp.float32
+
+    def setup(self):
+        self.kernel_net = [
+            SirenLayer(
+                output_dim=self.hidden_dim,
+                omega_0=self.omega_0,
+                is_first_layer=True,
+                dtype=self.dtype
+            )
+        ] + [
+            SirenLayer(
+                output_dim=self.hidden_dim,
+                omega_0=self.omega_0,
+                dtype=self.dtype
+            )
+            for _ in range(self.num_layers)
+        ]
+
+        self.output_linear = nn.Dense(
+            features=self.output_dim,
+            use_bias=True,
+            kernel_init=custom_uniform(numerator=1, mode="fan_in", distribution="normal", dtype=self.dtype),
+            bias_init=nn.initializers.zeros,
+            param_dtype=self.dtype
+        )
+
+    def __call__(self, x):
+        for layer in self.kernel_net:
+            x = layer(x)
+
+        out = self.output_linear(x)
+
+        return out
+
+
+class SirenLayer(nn.Module):
+    output_dim: int
+    omega_0: float
+    is_first_layer: bool = False
+    dtype: jnp.dtype = jnp.float32
+
+    def setup(self):
+        c = 1 if self.is_first_layer else 6 / self.omega_0**2
+        distrib = "uniform_squared" if self.is_first_layer else "uniform"
+        self.linear = nn.Dense(
+            features=self.output_dim,
+            use_bias=True,
+            kernel_init=custom_uniform(numerator=c, mode="fan_in", distribution=distrib, dtype=self.dtype),
+            bias_init=nn.initializers.zeros,
+            param_dtype=self.dtype
+        )
+
+    def __call__(self, x):
+        after_linear = self.omega_0 * self.linear(x)
+        return jnp.sin(after_linear)

utils.py

@@ -0,0 +1,167 @@
+import jax
+import math
+from typing import Any, Dict, Sequence, Union
+
+import jax.numpy as jnp
+from jax import dtypes, random
+from jax.nn.initializers import Initializer
+from typing import Callable
+from flax import linen as nn
+
+class FourierEmbs(nn.Module):
+    embed_scale: float
+    embed_dim: int
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(self, x):
+        kernel = self.param(
+            "kernel", jax.nn.initializers.normal(self.embed_scale, dtype=self.dtype), (x.shape[-1], self.embed_dim // 2)
+        )
+        y = jnp.concatenate(
+            [jnp.cos(jnp.dot(x, kernel)), jnp.sin(jnp.dot(x, kernel))], axis=-1
+        )
+        return y
+
+def _weight_fact(init_fn, mean, stddev, dtype=jnp.float32):
+    def init(key, shape):
+        key1, key2 = jax.random.split(key)
+        w = init_fn(key1, shape)
+        g = mean + nn.initializers.normal(stddev, dtype=dtype)(key2, (shape[-1],))
+        g = jnp.exp(g)
+        v = w / g
+        return g, v
+
+    return init
+
+class Dense(nn.Module):
+    features: int
+    kernel_init: Callable = nn.initializers.glorot_normal()
+    bias_init: Callable = nn.initializers.zeros
+    reparam : Union[None, Dict] = None
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(self, x):
+        if self.reparam is None:
+            kernel = self.param(
+                "kernel", self.kernel_init(dtype=self.dtype), (x.shape[-1], self.features)
+            )
+        elif self.reparam["type"] == "weight_fact":
+            g, v = self.param(
+                "kernel",
+                _weight_fact(
+                    self.kernel_init,
+                    mean=self.reparam["mean"],
+                    stddev=self.reparam["stddev"],
+                    dtype=self.dtype
+                ),
+                (x.shape[-1], self.features),
+            )
+
+            kernel = g * v
+
+        bias = self.param("bias", self.bias_init(dtype=self.dtype), (self.features,))
+
+        y = jnp.dot(x, kernel) + bias
+
+        return y
+
+
+def _compute_fans(
+    shape: tuple,
+    in_axis: Union[int, Sequence[int]] = -2,
+    out_axis: Union[int, Sequence[int]] = -1,
+    batch_axis: Union[int, Sequence[int]] = (),
+):
+    """Compute effective input and output sizes for a linear or convolutional layer.
+
+    Axes not in in_axis, out_axis, or batch_axis are assumed to constitute the "receptive field" of
+    a convolution (kernel spatial dimensions).
+    """
+    if len(shape) <= 1:
+        raise ValueError(
+            f"Can't compute input and output sizes of a {shape.rank}"
+            "-dimensional weights tensor. Must be at least 2D."
+        )
+
+    if isinstance(in_axis, int):
+        in_size = shape[in_axis]
+    else:
+        in_size = math.prod([shape[i] for i in in_axis])
+    if isinstance(out_axis, int):
+        out_size = shape[out_axis]
+    else:
+        out_size = math.prod([shape[i] for i in out_axis])
+    if isinstance(batch_axis, int):
+        batch_size = shape[batch_axis]
+    else:
+        batch_size = math.prod([shape[i] for i in batch_axis])
+    receptive_field_size = math.prod(shape) / in_size / out_size / batch_size
+    fan_in = in_size * receptive_field_size
+    fan_out = out_size * receptive_field_size
+    return fan_in, fan_out
+
+
+def custom_uniform(
+    numerator: float = 6,
+    mode: str = "fan_in",
+    dtype: jnp.dtype = jnp.float32,
+    in_axis: Union[int, Sequence[int]] = -2,
+    out_axis: Union[int, Sequence[int]] = -1,
+    batch_axis: Sequence[int] = (),
+    distribution: str = "uniform",
+) -> Initializer:
+    """Builds an initializer that returns real uniformly-distributed random arrays.
+
+    :param numerator: the numerator of the range of the random distribution.
+    :type numerator: float
+    :param mode: the mode for computing the range of the random distribution.
+    :type mode: str
+    :param dtype: optional; the initializer's default dtype.
+    :type dtype: jnp.dtype
+    :param in_axis: the axis or axes that specify the input size.
+    :type in_axis: Union[int, Sequence[int]]
+    :param out_axis: the axis or axes that specify the output size.
+    :type out_axis: Union[int, Sequence[int]]
+    :param batch_axis: the axis or axes that specify the batch size.
+    :type batch_axis: Sequence[int]
+    :param distribution: the distribution of the random distribution.
+    :type distribution: str
+
+    :return: An initializer that returns arrays whose values are uniformly distributed in
+        the range ``[-range, range)``.
+    :rtype: Initializer
+    """
+
+    def init(key: jax.random.key, shape: tuple, dtype: Any = dtype) -> Any:
+        dtype = dtypes.canonicalize_dtype(dtype)
+        fan_in, fan_out = _compute_fans(shape, in_axis, out_axis, batch_axis)
+        if mode == "fan_in":
+            denominator = fan_in
+        elif mode == "fan_out":
+            denominator = fan_out
+        elif mode == "fan_avg":
+            denominator = (fan_in + fan_out) / 2
+        else:
+            raise ValueError(f"invalid mode for variance scaling initializer: {mode}")
+        if distribution == "uniform":
+            return random.uniform(
+                key,
+                shape,
+                dtype,
+                minval=-jnp.sqrt(numerator / denominator),
+                maxval=jnp.sqrt(numerator / denominator),
+            )
+        elif distribution == "normal":
+            return random.normal(key, shape, dtype) * jnp.sqrt(numerator / denominator)
+        elif distribution == "uniform_squared":
+            return random.uniform(
+                key, shape, dtype, minval=-numerator / denominator, maxval=numerator / denominator
+            )
+        else:
+            raise ValueError(
+                f"invalid distribution for variance scaling initializer: {distribution}"
+            )
+
+    return init

wire.py

@@ -0,0 +1,135 @@
+import jax.numpy as jnp
+from flax import linen as nn
+from typing import Any
+
+import jax
+from .utils import custom_uniform
+from jax.nn.initializers import Initializer
+    
+
+def complex_kernel_uniform_init(numerator : float = 6,
+                                 mode : str = "fan_in",
+                                dtype : jnp.dtype = jnp.float32,
+                                distribution: str = "uniform") -> Initializer:
+    def init(key: jax.random.key, shape: tuple, dtype: Any = dtype) -> Any:
+        real_kernel = custom_uniform(numerator=numerator, mode=mode, distribution=distribution)(key, shape, dtype)
+        imag_kernel = custom_uniform(numerator=numerator, mode=mode, distribution=distribution)(key, shape, dtype)
+
+        return real_kernel + 1j * imag_kernel
+        
+    return init
+
+
+class WIRE(nn.Module):
+    output_dim: int
+    hidden_dim: int
+    num_layers: int
+    hidden_omega_0: float
+    first_omega_0: float
+    scale: float
+    complexgabor: bool = False
+    dtype: jnp.dtype = jnp.float32
+
+    def setup(self):
+        if self.complexgabor:
+            WIRElayer = ComplexGaborLayer
+            dtype = jnp.complex64
+        else:
+            WIRElayer = RealGaborLayer
+            dtype = self.dtype
+        self.kernel_net = [
+            WIRElayer(
+                output_dim=self.hidden_dim,
+                omega_0=self.first_omega_0,
+                s_0=self.scale,
+                is_first_layer=True,
+                dtype=dtype
+            )
+        ] + [
+            WIRElayer(
+                output_dim=self.hidden_dim,
+                omega_0=self.hidden_omega_0,
+                s_0=self.scale,
+                is_first_layer=False,
+                dtype=dtype
+            )
+            for _ in range(self.num_layers)
+        ]
+
+        self.output_linear = nn.Dense(
+            features=self.output_dim,
+            use_bias=True,
+            kernel_init=custom_uniform(numerator=1, mode="fan_in", distribution="normal"),
+            param_dtype=self.dtype,
+        )
+
+    def __call__(self, x):
+        for layer in self.kernel_net:
+            x = layer(x)
+
+        out = jnp.real(self.output_linear(x))
+
+        return out
+
+
+class ComplexGaborLayer(nn.Module):
+    output_dim: int
+    omega_0: float
+    s_0: float
+    is_first_layer: bool = False
+    dtype: jnp.dtype = jnp.float32
+
+    def setup(self):
+        c = 1 if self.is_first_layer else 6 / self.omega_0**2
+        distrib = "uniform_squared" if self.is_first_layer else "uniform"
+
+        if self.is_first_layer:
+            dtype = self.dtype
+        else:
+            dtype = jnp.complex64
+
+        self.linear = nn.Dense(
+            features=self.output_dim,
+            use_bias=True,
+            kernel_init=complex_kernel_uniform_init(numerator=c, mode="fan_in", distribution=distrib),
+            param_dtype=dtype
+        )
+
+    def __call__(self, x):
+        omega = self.omega_0 * self.linear(x)
+        scale = self.s_0 * self.linear(x)
+
+        return jnp.exp(1j * omega - (jnp.abs(scale)**2))
+
+
+class RealGaborLayer(nn.Module):
+    output_dim: int
+    omega_0: float
+    s_0: float
+    is_first_layer: bool = False
+    dtype: jnp.dtype = jnp.float32
+
+    def setup(self):
+
+        c = 1 if self.is_first_layer else 6 / self.omega_0**2
+        distrib = "uniform_squared" if self.is_first_layer else "uniform"
+
+        self.freqs = nn.Dense(
+            features=self.output_dim,
+            kernel_init=custom_uniform(numerator=c, mode="fan_in", distribution=distrib, dtype=self.dtype),
+            use_bias=True,
+            param_dtype=self.dtype
+        )
+
+        self.scales = nn.Dense(
+            features = self.output_dim,
+            kernel_init=custom_uniform(numerator=c, mode="fan_in", distribution=distrib, dtype=self.dtype),
+            use_bias=True,
+            param_dtype=self.dtype
+        )
+
+    def __call__(self, x):
+        omega = self.omega_0 * self.freqs(x)
+        scale = self.s_0 * self.scales(x)
+
+        return jnp.cos(omega) * jnp.exp(-(scale**2))