neuralop.training.tensorgrad.fno_tensorgrad_param_groups

neuralop.training.tensorgrad.fno_tensorgrad_param_groups(model: Module, rank: float = 0.25, sparse_ratio: float | None = None, min_params: int = 1000, update_proj_gap: int = 50, scale: float = 1.0, sparse_scale: float = 1.0, sparse_type: str = 'randk', lambda_sparse: float = 1.0, warm_restart: bool = True)[source]

Return AdamW and TensorGRaD parameter groups for FNO spectral weights.

The returned groups put large parameters from model.fno_blocks.convs in a TensorGRaD group and all remaining trainable parameters in a regular AdamW-style group.

Parameters:

modelnn.Module

FNO-like model with fno_blocks.convs.

rankfloat, optional

Tucker low-rank budget for selected spectral-convolution parameters.

sparse_ratiofloat, optional

Sparse residual budget. If None, only the low-rank branch is used.

min_paramsint, optional

Minimum number of entries a spectral-convolution parameter must have to be optimized with TensorGRaD.

update_proj_gapint, optional

Number of optimizer steps between projection updates.

scalefloat, optional

Scalar applied to projected-back low-rank updates.

sparse_scalefloat, optional

Scalar applied to projected-back sparse updates.

sparse_type{“randk”, “randomk”, “topk”}, optional

Strategy used by the sparse residual projector.

lambda_sparsefloat, optional

Weight used when adding the sparse update to the low-rank update.

warm_restartbool, optional

Whether Tucker decomposition reuses the previous projection factors when updating the low-rank basis.