Function tensor_plan 

pub fn tensor_plan(
    name: &str,
    mixer: MixerKind,
    _mlp: MlpKind,
    config: &ModelConfig,
) -> (TpShardKind, usize, usize)

(kind, full_out, full_in) for one checkpoint key.