param-lens

Install

No Marketplace listing yet; install the packaged .vsix from a release:

code --install-extension param-lens-0.1.0.vsix

What it does

Point it at a selection or a whole file: it looks for assignments that call a recognized nn.* layer — Linear, Conv2d, Embedding, LayerNorm, MultiheadAttention, LSTM, GRU — or another class defined in the same file, works out the parameter count from each call's arguments, and shows a breakdown grouped by class. It's a heuristic reader, regex plus light AST, not a Python interpreter: it never imports or runs the file, and when an argument can't be traced back to a literal number the line is listed as unresolved instead of guessed at.

In action

Param Lens: Count parameters in selection/file  (test/fixtures/nanogpt.py)
Total: 59,328 parameters

CausalSelfAttention (nn.Module) [used by: Block] — 4,224 params — line 19
  c_attn     nn.Linear     32x96 + 96 (bias)                 =  3,168
  c_proj     nn.Linear     32x32 + 32 (bias)                 =  1,056

GPT (nn.Module) [root] — 59,328 params — line 53
  wte        nn.Embedding  100x32                            =  3,200
  wpe        nn.Embedding  64x32                             =  2,048
  h          Block         Block() = 12,704 params/instance  = 50,816  x4
  ln_f       nn.LayerNorm  32 (weight) + 32 (bias)            =     64
  lm_head    nn.Linear     32x100 (bias=False)                =  3,200

Features

• —Recognizes nn.Linear, Conv2d, Embedding, LayerNorm, MultiheadAttention, LSTM and GRU.
• —Follows ModuleList([...]) and for-loops with append() to multiply repeated layers.
• —Resolves identifiers against literal values defined in the file, even config.n_embd.
• —An optional CodeLens shows ~N params above every class X(nn.Module).

Other tools