Model layers

PhiJAX provides coordinate feature mappings and dense-layer parameterizations used by its built-in architectures. They can also be composed into custom Flax NNX models.

Periodic features

PeriodicFeatures replaces each selected scalar coordinate with cosine and sine features, increasing the width by one per selected axis. Configured frequencies are angular frequencies, so the physical period is 2*pi/abs(frequency).

PeriodicFeatures [source]

PeriodicFeatures(
    input_dim: int,
    axes: Sequence[int],
    frequencies: Sequence[float],
)

Bases: Module

Replace selected coordinate axes with fixed cosine and sine features.

Attributes:

Name	Type	Description
`output_dim`		Embedded width, equal to `input_dim + len(axes)`.

Initialize deterministic periodic coordinate features.

Each configured coordinate x_i is replaced by (cos(w_i*x_i), sin(w_i*x_i)), so its physical period is 2*pi/abs(w_i).

Parameters:

Name	Type	Description	Default
`input_dim`	`int`	Width of the physical input coordinate vector.	required
`axes`	`Sequence[int]`	Unique coordinate axes to replace.	required
`frequencies`	`Sequence[float]`	Nonzero finite angular frequencies paired with `axes`.	required

Raises:

Type	Description
`ValueError`	If dimensions, axes, or frequencies are invalid.

call [source]

__call__(inputs: Array) -> jax.Array

Embed configured axes while preserving leading input dimensions.

Parameters:

Name	Type	Description	Default
`inputs`	`Array`	Coordinate array with final width :attr:`input_dim`.	required

Returns:

Type	Description
`Array`	Embedded array with final width :attr:`output_dim`.

Raises:

Type	Description
`ValueError`	If the final input width does not match :attr:`input_dim`.

Random Fourier features

RandomFourierFeatures returns 2 * embed_dim features using the paired mapping x -> [cos(xB), sin(xB)]. The construction follows the random-feature foundation of Rahimi and Recht (2007) and its use for coordinate networks by Tancik et al. (2020). PhiJAX stores the initially random projection B as a trainable parameter. Keeping B fixed would recover the classical random-feature interpretation.

RandomFourierFeatures [source]

RandomFourierFeatures(
    input_dim: int,
    embed_dim: int,
    *,
    scale: float = 1.0,
    compute_dtype: Any = jnp.float32,
    parameter_dtype: Any = jnp.float32,
    rngs: Rngs,
)

Bases: Module

Map coordinates to paired cosine and sine Fourier features through a trainable random projection.

The mapping x -> [cos(xB), sin(xB)] is related to random Fourier features for approximating shift-invariant kernels and to Fourier feature mappings for reducing the spectral bias of coordinate networks. PhiJAX samples B from a Gaussian distribution and stores it as a trainable :class:nnx.Param; classical random Fourier features instead keep the sampled projection fixed.

Attributes:

Name	Type	Description
`kernel`		Trainable Gaussian frequency matrix with shape `[input_dim, embed_dim]`.

References

Rahimi, A. and Recht, B. (2007). Random Features for Large-Scale Kernel Machines. NeurIPS. Tancik, M. et al. (2020). Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains. NeurIPS.

Initialize a trainable Gaussian Fourier kernel.

Parameters:

Name	Type	Description	Default
`input_dim`	`int`	Number of input coordinate channels.	required
`embed_dim`	`int`	Number of sampled frequencies. The output width is twice this value.	required
`scale`	`float`	Standard deviation of the sampled frequencies.	`1.0`
`compute_dtype`	`Any`	Data type used for projection arithmetic.	`float32`
`parameter_dtype`	`Any`	Data type used to store the trainable frequency matrix.	`float32`
`rngs`	`Rngs`	NNX random streams used for initialization.	required

Raises:

Type	Description
`ValueError`	If `input_dim` or `embed_dim` is not positive.

call [source]

__call__(inputs: Array) -> jax.Array

Embed inputs through paired cosine and sine projections.

Parameters:

Name	Type	Description	Default
`inputs`	`Array`	Coordinate array with final dimension `input_dim`.	required

Returns:

Type	Description
`Array`	Embedded array with unchanged leading dimensions and final width `2 * embed_dim`.

Factorized dense layers

FactorizedDense uses learned scale and direction parameters following Wang et al. (2022). PhiJAX stores dense weights in [input_dim, output_dim] layout, so its factorization is W = V diag(g): each entry of g scales one output column. With the common [output_dim, input_dim] layout, the same factorization is written W = diag(g) V.

FactorizedDense [source]

FactorizedDense(
    input_dim: int,
    output_dim: int,
    *,
    mean: float = 0.5,
    std: float = 0.1,
    use_bias: bool = True,
    kernel_init: Initializer | None = None,
    compute_dtype: Any = jnp.float32,
    parameter_dtype: Any = jnp.float32,
    rngs: Rngs,
)

Bases: Module

Apply a random weight-factorized affine transform, W = V diag(g).

Each output channel has a learned positive scale g_i and direction vector v_i. With PhiJAX's dense weight layout [input_dim, output_dim], scaling output columns gives W = V diag(g). This is the transpose-layout equivalent of the commonly written W = diag(g) V for [output_dim, input_dim] weights.

Attributes:

Name	Type	Description
`g`		Trainable positive scale vector with shape `[output_dim]`.
`v`		Trainable direction matrix with shape `[input_dim, output_dim]`.
`bias`		Optional trainable additive bias with shape `[output_dim]`.

References

Wang, S., Wang, H., Seidman, J. H., and Perdikaris, P. (2022). Random Weight Factorization Improves the Training of Continuous Neural Representations. arXiv:2210.01274.

Initialize factorized parameters with a reconstructed Xavier-uniform weight.

Parameters:

Name	Type	Description	Default
`input_dim`	`int`	Input feature width.	required
`output_dim`	`int`	Output feature width.	required
`mean`	`float`	Mean used to sample the log scale.	`0.5`
`std`	`float`	Standard deviation used to sample the log scale.	`0.1`
`use_bias`	`bool`	Whether to learn an additive bias.	`True`
`kernel_init`	`Initializer \| None`	Initializer for the reconstructed base weight. Defaults to Xavier uniform.	`None`
`compute_dtype`	`Any`	Data type used for matrix multiplication.	`float32`
`parameter_dtype`	`Any`	Data type used to store trainable parameters.	`float32`
`rngs`	`Rngs`	NNX random streams used for initialization.	required

Raises:

Type	Description
`ValueError`	If `input_dim` or `output_dim` is not positive.

weight `property`

weight: Array

Return the reconstructed dense weight.

Returns:

Type	Description
`Array`	Weight array in `[input_dim, output_dim]` layout.

call [source]

__call__(inputs: Array) -> jax.Array

Apply the factorized affine transform.

Parameters:

Name	Type	Description	Default
`inputs`	`Array`	Array whose final dimension is `input_dim`.	required

Returns:

Type	Description
`Array`	Array with unchanged leading dimensions and final width `output_dim`.

Model layers

Periodic features

PeriodicFeatures [source]

__call__ [source]

Random Fourier features

RandomFourierFeatures [source]

__call__ [source]

Factorized dense layers

FactorizedDense [source]

weight property

__call__ [source]

call [source]

call [source]

weight `property`

call [source]