Unsupported tensor rank for optimizer state: {other}

Error message

Unsupported tensor rank for optimizer state: {other}

What it means

burn-optim's optimizer state machinery is generated by a macro that only instantiates tensor ranks 1 through 8; any other rank hits the `other =>` catch-all and panics with 'Unsupported tensor rank for optimizer state'. The runtime rank of a parameter tensor is dispatched at runtime because the optimizer state view is non-generic over rank.

Source

Thrown at crates/burn-optim/src/optim/module/base.rs:126

                $body
            }
            5 => {
                const $d: usize = 5;
                $body
            }
            6 => {
                const $d: usize = 6;
                $body
            }
            7 => {
                const $d: usize = 7;
                $body
            }
            8 => {
                const $d: usize = 8;
                $body
            }
            other => panic!("Unsupported tensor rank for optimizer state: {other}"),
        }
    };
}

/// Object-safe view over an [`Optimizer`], allowing [`ModuleOptimizer`](crate::optim::ModuleOptimizer)
/// to stay non-generic. Rank-generic operations are dispatched on a runtime rank.
pub trait DynOptimizer: Send + Sync {
    /// Perform an optimizer step for a single parameter of the given `rank`.
    fn step_dyn(
        &self,
        rank: usize,
        lr: LearningRate,
        tensor: BridgeTensor,
        grad: BridgeTensor,
        state: Option<DynState>,
    ) -> (BridgeTensor, Option<DynState>);

    /// Move a state to the given device.

View on GitHub (pinned to d16f7ba2ed)

Solutions

Reshape the parameter to rank between 1 and 8 (e.g. wrap scalars as rank-1 tensors of length 1).
If rank > 8 is truly required, extend the rank-dispatch macro arms to cover the needed rank.
Split overly large multi-dimensional parameters into fewer-dimensional components.

Example fix

// before
let w: Tensor<B, 1> = Tensor::from_floats([0.5]); // scalar-like state, ok
// panic case: rank 9 tensor used as a parameter
// after
let w = big_tensor.reshape([d1, d2, d3, d4, d5, d6, d7, d8]); // rank <= 8

Defensive patterns

Strategy: validation

Validate before calling

fn optimizer_state_rank_ok(param_rank: usize) -> bool {
    (1..=8).contains(&param_rank) // macro only instantiates ranks 1..=8
}

Try / catch

// validate parameter shapes before step()/checkpointing
for param in module.parameters() {
    let rank = param.shape().dims().len();
    assert!((1..=8).contains(&rank), "parameter rank {rank} unsupported by optimizer state");
}

Prevention

Keep all trainable parameters at rank 1..=8; wrap scalars as rank-1 tensors.
Avoid >8-dimensional parameters; split them into multiple lower-rank parameters.
Add a startup pass that asserts every module parameter's rank is within 1..=8 before optimizer init.

When it happens

Trigger: Running an optimizer (e.g. Adam via ModuleOptimizer) whose adaptive state must be allocated for a parameter tensor with rank 0 (scalar) or rank > 8.

Common situations: Optimizing a scalar parameter (rank 0); extreme models with deeply nested, >8-dimensional tensors; custom modules exposing unusual parameter shapes into the optimizer state checkpoint path (burn-autodiff CheckpointerBuilder::extend flows state through this dispatch).

Related errors

AI-assisted analysis of tracel-ai/burn@d16f7ba2ed (2026-09-05). Data as JSON: /api/errors/c968c26ffb0edade. Report an issue: GitHub.