Adding diagrams to my static site generator with D2

A lot of the time when I've been writing posts for this blog, I've felt that a diagram would really help. But they're a pain to produce well, and I think I underuse them as a result. I wanted to fix that, and wound up adding D2 support to my static site generator. I think it works pretty well!

In the past, I've tried drawing my own diagrams in LibreOffice and exporting as SVG, but my complete lack of artistic skill doesn't help:

Asking an AI to do it for me helped in simple cases:

...but with something less standard (there must be a million neural network diagrams in their training sets) it can be really fiddly to get something right.

I did some investigations into the various diagram-generating tools out there, and decided to give D2 a go. It has a simple language for specifying what your diagram should show, and the output is pretty nice:

Here's the source for that diagram:

style.fill: transparent

tokens: Tokens
tokens -> llm.token-embeddings

llm: "" {
  token-embeddings: Token embeddings
  position-embeddings: Position embeddings

  plus: "+" {
    shape: circle
    width: 36
    height: 36
    style.font-size: 24
    style.fill: transparent
  }

  token-embeddings -> plus
  position-embeddings -> plus

  input-embeddings: Input embeddings
  plus -> input-embeddings

  transformers-layers: "" {
    style.stroke-dash: 3
    style.fill: transparent
    transformers-1: Transformers layer 1
    transformers-2: Transformers layer 2
    dots: "⋮" {shape: text; style.font-size: 28}
    transformers-n: Transformers layer n
    transformers-1 -> transformers-2 -> dots -> transformers-n
  }
  input-embeddings -> transformers-layers.transformers-1

  final-norm: LayerNorm
  transformers-layers.transformers-n -> final-norm

  output-head: Output head
  final-norm -> output-head
}

output-logits: Logits
llm.output-head -> output-logits

That looks pretty clear to me!

So now, in the source for my blog posts, I have a diagrams directory. That contains subdirectories -- by convention, I create one for each post that needs diagrams -- and D2 files. These can be generated automatically when I publish:

    compile_d2_diagrams(input_path=Path(INPUT_DIR) / "diagrams", output_path=Path(NEW_OUTPUT_DIR) / "diagrams")

...

def compile_d2_diagrams(input_path, output_path):
    if input_path.is_file():
        if not input_path.name.endswith(".d2"):
            raise Exception(f"Unknown file type in D2 tree: {input_path}")
        output_path = output_path.with_suffix(".svg")
        print(f"Compiling D2 diagram in {input_path} to {output_path}")
        subprocess.check_call([
            "d2",
            "--pad=0",
            "--layout=elk",
            "--elk-nodeNodeBetweenLayers=30",
            "--elk-padding=[top=20,left=20,bottom=20,right=20]",
            input_path,
            output_path
        ])
        return
    if input_path.is_dir():
        print(f"Making diagram directory to match {input_path}: {output_path}")
        output_path.mkdir()
        for child_path in input_path.iterdir():
            compile_d2_diagrams(child_path, output_path / child_path.name)
        return
    raise Exception(f"Don't know what {input_path} is!")

(Hat tip to Evan Hahn for the with_suffix method on Path, which I wasn't aware of.)

The flags on the command line took a little bit of fiddling; the --pad=0 just gets rid of the large margins that D2 puts around the diagram by default, but the others are to tell it to use the ELK layout package with particular formatting. Its default layout has curvy lines, and I prefer the closer-to-right-angle ones that ELK provides.

Another awkward bit was in scaling; the file that is generated by that d2 command comes out pretty large (you can see it full-size here). By default, I allow images inlined into my posts to be as wide as the text, but that would still be too large here.

I use markdown2 to convert the markdown source for my posts into HTML, and there isn't any way to tell it what size an image should be using markdown-ish syntax. So for now, instead of embedding images the normal markdown way, like this:

![A simple neural network](/post-assets/neural-networks-maths/network.svg "A simple neural network")

...for these D2-generated ones I'll just embed a normal tag like this:

...so that I can control the size. Perhaps more work needed there.

At some point I may go back and update my old diagrams -- at least, the really ugly hand-drawn ones -- to use this.

And a random thought: perhaps it might also make sense to include the D2 source somehow on the blog? I can imagine that it could help with accessibility in some situations, and perhaps also for any LLMs stopping by. Will have to ponder that a bit more.

What do you think? Does the D2 diagram look good to you? Or is there a better diagramming package that might work better?