TransformerBlock · 1.77M params · 11.6GFLOP · 215 ops → 95 fused · CPU/XLA
input (f32(8, 128, 384), f32(8, 128, 384))

memory
total         █████████████████████████  64.5 MB
params        ██▋······················  6.76 MB
input         ···▉·····················  3 MB
intermediate  ····██████████████████▌··  48 MB
output        ·······················██  6.76 MB

layer                             FLOPs              share  ops  output
TransformerBlock                  11.6G  ██████████ 100.0%  215  f32(8, 128, 384)
  attention MultiHeadAttention    4.25G  ███▊······  36.8%   87  f32(8, 128, 384)
  attention_norm LayerNorm        2.76M  ··········   0.0%   23  f32(8, 128, 384)
  mlp_norm LayerNorm               7.5M  ··········   0.1%   49  f32(8, 128, 384)
  mlp_in Linear                   3.62G  ····██▉···  31.4%    4  f32(8, 128, 1536)
  mlp_out Linear                  3.62G  ·······███  31.4%    4  f32(8, 128, 384)