import jax
import jax.numpy as jnp
from flax import nnx
from typing import Union
from .utils import solve_full_lstsq
# Dictionary of Legendre polynomials up to degree 20
Lg = {
0: lambda x: jnp.ones_like(x),
1: lambda x: x,
2: lambda x: (1/2) * (3 * x**2 - 1),
3: lambda x: (1/2) * (5 * x**3 - 3 * x),
4: lambda x: (1/8) * (35 * x**4 - 30 * x**2 + 3),
5: lambda x: (1/8) * (63 * x**5 - 70 * x**3 + 15 * x),
6: lambda x: (1/16) * (231 * x**6 - 315 * x**4 + 105 * x**2 - 5),
7: lambda x: (1/16) * (429 * x**7 - 693 * x**5 + 315 * x**3 - 35*x),
8: lambda x: (1/128) * (6435 * x**8 - 12012 * x**6 + 6930 * x**4 - 1260 * x**2 + 35),
9: lambda x: (1/128) * (12155 * x**9 - 25740 * x**7 + 18018 * x**5 - 4620 * x**3 + 315 * x),
10: lambda x: (1/256) * (46189 * x**10 - 109395 * x**8 + 90090 * x**6 - 30030 * x**4 + 3465 * x**2 - 63),
11: lambda x: (1/256) * (88179 * x**11 - 230945 * x**9 + 218790 * x**7 - 90090 * x**5 + 15015 * x**3 - 693 * x),
12: lambda x: (1/1024) * (676039 * x**12 - 1939938 * x**10 + 2078505 * x**8 - 1021020 * x**6 + 225225 * x**4 - 18018 * x**2 + 231),
13: lambda x: (1/1024) * (1300075 * x**13 - 4056234 * x**11 + 4849845 * x**9 - 2771340 * x**7 + 765765 * x**5 - 90090 * x**3 + 3003 * x),
14: lambda x: (1/2048) * (5014575 * x**14 - 16900975 * x**12 + 22309287 * x**10 - 14549535 * x**8 + 4849845 * x**6 - 765765 * x**4 + 45045 * x**2 - 429),
15: lambda x: (1/2048) * (9694845 * x**15 - 35102025 * x**13 + 50702925 * x**11 - 37182145 * x**9 + 14549535 * x**7 - 2909907 * x**5 + 255255 * x**3 - 6435 * x),
16: lambda x: (1/32768) * (300540195 * x**16 - 1163381400 * x**14 + 1825305300 * x**12 - 1487285800 * x**10 + 669278610 * x**8 - 162954792 * x**6 + 19399380 * x**4 - 875160 * x**2 + 6435),
17: lambda x: (1/32768) * (583401555 * x**17 - 2404321560 * x**15 + 4071834900 * x**13 - 3650610600 * x**11 + 1859107250 * x**9 - 535422888 * x**7 + 81477396 * x**5 - 5542680 * x**3 + 109395 * x),
18: lambda x: (1/65536) * (2268783825 * x**18 - 9917826435 * x**16 + 18032411700 * x**14 - 17644617900 * x**12 + 10039179150 * x**10 - 3346393050 * x**8 + 624660036 * x**6 - 58198140 * x**4 + 2078505 * x**2 - 12155),
19: lambda x: (1/65536) * (4418157975 * x**19 - 20419054425 * x**17 + 39671305740 * x**15 - 42075627300 * x**13 + 26466926850 * x**11 - 10039179150 * x**9 + 2230928700 * x**7 - 267711444 * x**5 + 14549535 * x**3 - 230945 * x),
20: lambda x: (1/262144) * (34461632205 * x**20 - 167890003050 * x**18 + 347123925225 * x**16 - 396713057400 * x**14 + 273491577450 * x**12 - 116454478140 * x**10 + 30117537450 * x**8 - 4461857400 * x**6 + 334639305 * x**4 - 9699690 * x**2 + 46189),
}
[docs]
class LegendreLayer(nnx.Module):
"""
LegendreLayer class. Corresponds to the Legendre version of KANs and comes in two "flavors":
"default": uses the recursion formula to calculate polynomials up to arbitrary degree
"exact": uses pre-defined functions for higher efficiency, but cannot scale up to arbitrary degrees
Attributes:
n_in (int):
Number of layer's incoming nodes.
n_out (int):
Number of layer's outgoing nodes.
D (int):
Degree of Legendre polynomial.
flavor (str):
One of "default" or "exact" - chooses basis implementation.
residual (Union[nnx.Module, None]):
Function that is applied on samples to calculate residual activation.
rngs (nnx.Rngs):
Random number generator state.
bias (Union[nnx.Param, None]):
Bias parameter if add_bias is True, else None.
c_ext (Union[nnx.Param, None]):
External weights if external_weights is True, else None.
c_basis (nnx.Param):
Trainable coefficients for the basis functions.
c_res (Union[nnx.Param, None]):
Trainable coefficients for residual activation if residual is not None.
"""
[docs]
def __init__(self, n_in: int = 2, n_out: int = 5, D: int = 5, flavor: Union[str, None] = None,
residual: Union[nnx.Module, None] = None, external_weights: bool = False,
init_scheme: Union[dict, None] = None, add_bias: bool = True, seed: int = 42):
"""
Initializes a LegendreLayer instance.
Args:
n_in (int):
Number of layer's incoming nodes.
n_out (int):
Number of layer's outgoing nodes.
D (int):
Degree of Legendre polynomial.
flavor (Union[str, None]):
One of "default", "modified", or "exact" - chooses basis implementation.
residual (Union[nnx.Module, None]):
Function that is applied on samples to calculate residual activation.
external_weights (bool):
Boolean that controls if the trainable weights (n_out, n_in) should be applied to the activations.
init_scheme (Union[dict, None]):
Dictionary that defines how the trainable parameters of the layer are initialized.
add_bias (bool):
Boolean that controls wether bias terms are also included during the forward pass or not.
seed (int):
Random key selection for initializations wherever necessary.
Example:
>>> layer = LegendreLayer(n_in = 2, n_out = 5, D = 5, flavor = "default",
>>> residual = None, external_weights = False, init_scheme = None,
>>> add_bias = True, seed = 42)
"""
if flavor is None:
flavor = "default"
elif flavor == "exact":
max_deg = max(list(Lg.keys()))
if D > max_deg:
raise ValueError(f"For method 'exact', the maximum degree cannot exceed {max_deg}.")
# Setup basic parameters
self.n_in = n_in
self.n_out = n_out
self.D = D
self.flavor = flavor
self.residual = residual
# Setup nnx rngs
self.rngs = nnx.Rngs(seed)
# Add bias
if add_bias == True:
self.bias = nnx.Param(jnp.zeros((n_out,)))
else:
self.bias = None
# If external_weights == True, we initialize weights for the activation functions equal to unity
if external_weights == True:
self.c_ext = nnx.Param(
nnx.initializers.ones(
self.rngs.params(), (self.n_out, self.n_in), jnp.float32)
)
else:
self.c_ext = None
# Initialize the remaining trainable parameters, based on the selected initialization scheme
c_res, c_basis = self._initialize_params(init_scheme, seed)
self.c_basis = nnx.Param(c_basis)
if residual is not None:
self.c_res = nnx.Param(c_res)
[docs]
def basis(self, x):
"""
Based on the degree and flavor, the values of the Legendre basis functions are calculated on the input.
Args:
x (jnp.array):
Inputs, shape (batch, n_in).
Returns:
leg (jnp.array):
Legendre basis functions applied on inputs, shape (batch, n_in, D+1).
Example:
>>> layer = LegendreLayer(n_in = 2, n_out = 5, D = 5, flavor = "default",
>>> residual = None, external_weights = False, init_scheme = None,
>>> add_bias = True, seed = 42)
>>>
>>> key = jax.random.key(42)
>>> x_batch = jax.random.uniform(key, shape=(100, 2), minval=-1.0, maxval=1.0)
>>>
>>> output = layer.basis(x_batch)
"""
batch = x.shape[0]
# Apply tanh activation
x = jnp.tanh(x) # (batch, n_in)
# Default implementation
if self.flavor == "default":
# Order 0 is set by default, since we initialize at 1
leg = jnp.ones((batch, self.n_in, self.D+1))
# Set order 1 as well
leg = leg.at[:, :, 1].set(x)
# Handle higher orders iteratively
for n in range(1, self.D):
leg = leg.at[:, :, n+1].set(((2 * n + 1) * x * leg[:, :, n] - n * leg[:, :, n-1]) / (n + 1))
# Exact calculation of polynomials
elif self.flavor == "exact":
leg = jnp.stack([Lg[i](x) for i in range(self.D + 1)], axis=-1) # (batch, n_in, D+1)
# Other flavor
else:
raise ValueError(f"Unknown layer flavor: {self.flavor}")
# Exclude the constant "1" dimension if bias is included
if self.bias is not None:
return leg[:, :, 1:]
else:
return leg
def _initialize_params(self, init_scheme, seed):
"""
Initializes the c_res (if residual is activated) and c_basis trainable parameters (only used in __init__)
Args:
init_scheme (Union[dict, None]):
Dictionary that defines how the trainable parameters of the layer are initialized. Options: "default", "lecun", "custom"
"""
if init_scheme is None:
init_scheme = {"type" : "default"}
init_type = init_scheme.get("type", "default")
# Case where no residual is used
if self.residual is None:
c_res = None
# Find if we have D+1 external dimension (if add_bias = False) or D (if add_bias = True)
ext_dim = self.D if self.bias is not None else self.D+1
# Default initialization
if init_type == "default":
if self.residual is not None:
c_res = nnx.initializers.glorot_uniform(in_axis=-1, out_axis=-2)(
self.rngs.params(), (self.n_out, self.n_in), jnp.float32
)
std = 1.0/jnp.sqrt(self.n_in * ext_dim)
c_basis = nnx.initializers.truncated_normal(stddev=std)(
self.rngs.params(), (self.n_out, self.n_in, ext_dim), jnp.float32
)
# Custom power law initialization
# c_basis ~ N(0, s_b), s_b = const_b / [ (ext_dim+1)^pow_b1 * n_in^pow_b2 ]
# c_res ~ N(0, s_r), s_r = const_r / [ (ext_dim+1)^pow_r1 * n_in^pow_r2 ]
elif init_type == "power":
const_b = init_scheme.get("const_b", 1.0)
pow_b1 = init_scheme.get("pow_b1", 0.5)
pow_b2 = init_scheme.get("pow_b2", 0.5)
if self.residual is not None:
basis_term = ext_dim + 1
const_r = init_scheme.get("const_r", 1.0)
pow_r1 = init_scheme.get("pow_r1", 0.5)
pow_r2 = init_scheme.get("pow_r2", 0.5)
std_res = const_r / ( (basis_term**pow_r1) * (self.n_in**pow_r2) )
c_res = nnx.initializers.normal(stddev=std_res)(
self.rngs.params(), (self.n_out, self.n_in), jnp.float32
)
else:
basis_term = ext_dim
std_b = const_b / ( (basis_term**pow_b1) * (self.n_in**pow_b2) )
c_basis = nnx.initializers.normal(stddev=std_b)(
self.rngs.params(), (self.n_out, self.n_in, ext_dim), jnp.float32
)
# LeCun-like initialization, where Var[in] = Var[out]
elif init_type == "lecun":
key = jax.random.key(seed)
# Also get distribution type
distrib = init_scheme.get("distribution", "uniform")
if distrib is None:
distrib = "uniform"
sample_size = init_scheme.get("sample_size", 10000)
if sample_size is None:
sample_size = 10000
# Generate a sample of points
if distrib == "uniform":
sample = jax.random.uniform(key, shape=(sample_size,), minval=-1.0, maxval=1.0)
elif distrib == "normal":
sample = jax.random.normal(key, shape=(sample_size,))
# Finally get gain
gain = init_scheme.get("gain", None)
if gain is None:
gain = sample.std().item()
# Extend the sample to be able to pass through basis
sample_ext = jnp.tile(sample[:, None], (1, self.n_in))
# Calculate B_m^2(x)
y_b = self.basis(sample_ext)
# Calculate the average of B_m^2(x)
y_b_sq = y_b**2
y_b_sq_mean = y_b_sq.mean().item()
if self.residual is not None:
# Variance equipartitioned across all terms
scale = self.n_in * (ext_dim + 1)
# Apply the residual function
y_res = self.residual(sample)
# Calculate the average of residual^2(x)
y_res_sq = y_res**2
y_res_sq_mean = y_res_sq.mean().item()
std_res = gain/jnp.sqrt(scale*y_res_sq_mean)
c_res = nnx.initializers.normal(stddev=std_res)(self.rngs.params(), (self.n_out, self.n_in), jnp.float32)
else:
# Variance equipartitioned across G+k terms
scale = self.n_in * ext_dim
std_b = gain/jnp.sqrt(scale*y_b_sq_mean)
c_basis = nnx.initializers.normal(stddev=std_b)(
self.rngs.params(), (self.n_out, self.n_in, ext_dim), jnp.float32
)
# Glorot-like initialization, where we attempt to balance Var[in] = Var[out] and Var[δin] = Var[δout]
elif init_type == "glorot":
key = jax.random.key(seed)
# Also get distribution type
distrib = init_scheme.get("distribution", "uniform")
if distrib is None:
distrib = "uniform"
sample_size = init_scheme.get("sample_size", 10000)
if sample_size is None:
sample_size = 10000
# Generate a sample of points
if distrib == "uniform":
sample = jax.random.uniform(key, shape=(sample_size,), minval=-1.0, maxval=1.0)
elif distrib == "normal":
sample = jax.random.normal(key, shape=(sample_size,))
# Finally get gain
gain = init_scheme.get("gain", None)
if gain is None:
gain = sample.std().item()
# Extend the sample to be able to pass through basis
sample_ext = jnp.tile(sample[:, None], (1, self.n_in))
# ------------- Basis function gradient ----------------------
# Define a scalar version of the basis function
def basis_scalar(x):
return self.basis(jnp.array([[x]]))[0, 0, :]
# Create a Jacobian function for the scalar wrapper
jac_basis = jax.jacobian(basis_scalar)
num_batches = 20
batch_size = sample_size // num_batches
grad_sq_accum = 0.0
for i in range(num_batches):
batch = sample[i*batch_size:(i+1)*batch_size]
grad_batch = jax.vmap(jac_basis)(batch)
grad_sq_accum += (grad_batch**2).sum()
# Calculate E[B'_m^2(x)]
grad_b_sq_mean = grad_sq_accum / (sample_size * ext_dim)
# ------------------------------------------------------------
# Calculate E[B_m^2(x)]
y_b = self.basis(sample_ext)
y_b_sq = y_b**2
y_b_sq_mean = y_b_sq.mean().item()
# Deal with residual if available
if self.residual is not None:
# Variance equipartitioned across all terms
scale_in = self.n_in * (ext_dim + 1)
scale_out = self.n_out * (ext_dim + 1)
# ------------- Residual function gradient ----------------------
# Similar idea to the basis function
def r(x):
return self.residual(x)
jac_res = jax.jacobian(r)
grad_res = jax.vmap(jac_res)(sample)
# ------------------------------------------------------------
# Calculate E[R^2(x)]
y_res = self.residual(sample)
y_res_sq = y_res**2
y_res_sq_mean = y_res_sq.mean().item()
# Calculate E[R'^2(x)]
grad_res_sq = grad_res**2
grad_res_sq_mean = grad_res_sq.mean().item()
std_res = gain*jnp.sqrt(2.0 / (scale_in*y_res_sq_mean + scale_out*grad_res_sq_mean))
c_res = nnx.initializers.normal(stddev=std_res)(self.rngs.params(), (self.n_out, self.n_in), jnp.float32)
else:
# Variance equipartitioned across G+k terms
scale_in = self.n_in * ext_dim
scale_out = self.n_out * ext_dim
std_b = gain*jnp.sqrt(2.0 / (scale_in*y_b_sq_mean + scale_out*grad_b_sq_mean))
c_basis = nnx.initializers.normal(stddev=std_b)(
self.rngs.params(), (self.n_out, self.n_in, ext_dim), jnp.float32
)
# Glorot-like initialization as presented in the paper "Training Deep Physics-Informed Kolmogorov-Arnold Networks"
# https://www.sciencedirect.com/science/article/pii/S0045782526000356
# The main difference is that we do not aggregate over all sigmas, each mode has its own, hence "fine grained"
elif init_type == "glorot_fine":
key = jax.random.key(seed)
# Also get distribution type
distrib = init_scheme.get("distribution", "uniform")
if distrib is None:
distrib = "uniform"
sample_size = init_scheme.get("sample_size", 10000)
if sample_size is None:
sample_size = 10000
# Generate a sample of points
if distrib == "uniform":
sample = jax.random.uniform(key, shape=(sample_size,), minval=-1.0, maxval=1.0)
elif distrib == "normal":
sample = jax.random.normal(key, shape=(sample_size,))
# Finally get gain
gain = init_scheme.get("gain", None)
if gain is None:
gain = sample.std().item()
# Extend the sample to be able to pass through basis
sample_ext = jnp.tile(sample[:, None], (1, self.n_in))
# ------------- Basis functions ------------------------
# μ⁽0⁾ₘ (⟨B_m²⟩)
B = self.basis(sample_ext)
mu0 = (B**2).mean(axis=(0, 1))
# μ⁽1⁾ₘ (⟨B'_m²⟩)
# Define a scalar version of the basis function
basis_scalar = lambda x: self.basis(jnp.array([[x]]))[0, 0, :]
jac_basis = jax.jacrev(basis_scalar)
mu1 = (jax.vmap(jac_basis)(sample)**2).mean(axis=0)
# ------------- Residual function ----------------------
# Deal with residual if available - same as simple glorot
if self.residual is not None:
# Variance equipartitioned across all terms
scale_in = self.n_in * (ext_dim + 1)
scale_out = self.n_out * (ext_dim + 1)
# ------------- Residual function gradient ----------------------
# Similar idea to the basis function
def r(x):
return self.residual(x)
jac_res = jax.jacobian(r)
grad_res = jax.vmap(jac_res)(sample)
# ------------------------------------------------------------
# Calculate E[R^2(x)]
y_res = self.residual(sample)
y_res_sq = y_res**2
y_res_sq_mean = y_res_sq.mean().item()
# Calculate E[R'^2(x)]
grad_res_sq = grad_res**2
grad_res_sq_mean = grad_res_sq.mean().item()
std_res = gain*jnp.sqrt(2.0 / (scale_in*y_res_sq_mean + scale_out*grad_res_sq_mean))
c_res = nnx.initializers.normal(stddev=std_res)(self.rngs.params(), (self.n_out, self.n_in), jnp.float32)
else:
# Variance equipartitioned across G+k terms
scale_in = self.n_in * ext_dim
scale_out = self.n_out * ext_dim
sigma_vec = gain * jnp.sqrt(1.0 / (scale_in*mu0 + scale_out*mu1))
noise = nnx.initializers.normal(stddev=1.0)(
self.rngs.params(), (self.n_out, self.n_in, ext_dim), jnp.float32
)
c_basis = noise * sigma_vec
# Custom initialization, where the user inputs pre-determined arrays
elif init_type == "custom":
if self.residual is not None:
c_res = init_scheme.get("c_res", None)
c_basis = init_scheme.get("c_basis", None)
else:
raise ValueError(f"Unknown initialization method: {init_type}")
return c_res, c_basis
[docs]
def update_grid(self, x, D_new):
"""
For the case of LegendreKANs there is no concept of grid. However, a fine-graining approach can be followed by progressively increasing the degree of the polynomials.
Args:
x (jnp.array):
Inputs, shape (batch, n_in).
D_new (int):
New Legendre polynomial degree.
Example:
>>> layer = LegendreLayer(n_in = 2, n_out = 5, D = 5, flavor = "default",
>>> residual = None, external_weights = False, init_scheme = None,
>>> add_bias = True, seed = 42)
>>>
>>> key = jax.random.key(42)
>>> x_batch = jax.random.uniform(key, shape=(100, 2), minval=-1.0, maxval=1.0)
>>>
>>> layer.update_grid(x=x_batch, D_new=8)
"""
# Apply the inputs to the current grid to acquire y = Sum(ciBi(x)), where ci are
# the current coefficients and Bi(x) are the current Legendre basis functions
Bi = self.basis(x).transpose(1, 0, 2) # (n_in, batch, D+1)
ci = self.c_basis[...].transpose(1, 2, 0) # (n_in, D+1, n_out)
ciBi = jnp.einsum('ijk,ikm->ijm', Bi, ci) # (n_in, batch, n_out)
# Update the degree order
self.D = D_new
# Get the Bj(x) for the degree order
Bj = self.basis(x).transpose(1, 0, 2) # (n_in, batch, D_new+1)
# Solve for the new coefficients
cj = solve_full_lstsq(Bj, ciBi) # (n_in, D_new+1, n_out)
# Cast into shape (n_out, n_in, D_new+1)
cj = cj.transpose(2, 0, 1)
self.c_basis = nnx.Param(cj)
[docs]
def __call__(self, x):
"""
The layer's forward pass.
Args:
x (jnp.array):
Inputs, shape (batch, n_in).
Returns:
y (jnp.array):
Output of the forward pass, shape (batch, n_out).
Example:
>>> layer = LegendreLayer(n_in = 2, n_out = 5, D = 5, flavor = "default",
>>> residual = None, external_weights = False, init_scheme = None,
>>> add_bias = True, seed = 42)
>>>
>>> key = jax.random.key(42)
>>> x_batch = jax.random.uniform(key, shape=(100, 2), minval=-1.0, maxval=1.0)
>>>
>>> output = layer(x_batch)
"""
batch = x.shape[0]
# Calculate basis activations
Bi = self.basis(x) # (batch, n_in, D+1)
act = Bi.reshape(batch, -1) # (batch, n_in * (D+1))
# Check if external_weights == True
if self.c_ext is not None:
act_w = self.c_basis[...] * self.c_ext[..., None] # (n_out, n_in, D+1)
else:
act_w = self.c_basis[...]
# Calculate coefficients
act_w = act_w.reshape(self.n_out, -1) # (n_out, n_in * (D+1))
y = jnp.matmul(act, act_w.T) # (batch, n_out)
# Check if there is a residual function
if self.residual is not None:
# Calculate residual activation
res = self.residual(x) # (batch, n_in)
# Multiply by trainable weights
res_w = self.c_res[...] # (n_out, n_in)
full_res = jnp.matmul(res, res_w.T) # (batch, n_out)
y += full_res # (batch, n_out)
if self.bias is not None:
y += self.bias[...] # (batch, n_out)
return y