book_layout: Zero support, wide-diagram scaling, public API

Completing the engine:
- Zero renders as a labeled 0 node (a 0 scalar for order-0); previously
  raised NotImplementedError.
- to_book_tikz gains scale= and max_width=: a diagram wider than
  max_width is uniformly scaled (transform shape) to fit, so wide
  gradients/products stay on the page instead of overflowing.
- Public API: exported as tensorgrad.to_book_tikz and added a
  Tensor.to_book_tikz(**kwargs) convenience method (lazy import). All
  three call paths agree.
- Fixed a SyntaxWarning (module docstring is now a raw string; it
  contains \dloop).

31 tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018PzG3QbNtaFmABBHBG39wp

Author: thomasahle

tensorgrad/__init__.py

@@ -13,5 +13,6 @@
 from .typing import typed  # noqa: F401
 from .compiler.runtime import Output, compile, grad  # noqa: F401
 from .extras.expectation import Expectation  # noqa: F401
+from .extras.book_layout import to_book_tikz  # noqa: F401
 from .functions import frobenius2, kronecker, diag, sum, log, pow, trace  # noqa: F401
 from sympy import symbols  # noqa: F401

tensorgrad/extras/book_layout.py

@@ -1,4 +1,4 @@
-"""Book-grammar layout for tensor diagrams.
+r"""Book-grammar layout for tensor diagrams.
 
 Turns a tensorgrad ``Tensor`` into a diagram laid out the way the Tensor
 Cookbook draws them by hand.  This is deliberately NOT a general graph-drawing
@@ -51,7 +51,9 @@
 from numbers import Number
 from typing import Optional
 
-from tensorgrad.tensor import Delta, Derivative, Function, Product, Rename, Sum, Tensor, Variable
+from tensorgrad.tensor import (
+    Delta, Derivative, Function, Product, Rename, Sum, Tensor, Variable, Zero,
+)
 
 try:
     from tensorgrad.extras.expectation import Expectation
@@ -165,6 +167,15 @@ def walk(t: Tensor) -> dict[str, str]:
             label = t.name
             g.atoms.append(AtomSpec(aid, "var", label, toks))
             return dict(zip(list(t.edges), toks))
+        if isinstance(t, Zero):
+            edges = list(t.edges)
+            aid = len(g.atoms)
+            if not edges:
+                g.atoms.append(AtomSpec(aid, "scalar", "0", []))
+                return {}
+            toks = [fresh("w") for _ in edges]
+            g.atoms.append(AtomSpec(aid, "var", "0", toks))
+            return dict(zip(edges, toks))
         if isinstance(t, Delta):
             edges = list(t.edges)
             if len(edges) == 2:
@@ -1211,11 +1222,28 @@ def to_book_tikz(
     left: Optional[str] = None,
     right: Optional[str] = None,
     baseline: str = "-.25em",
+    scale: Optional[float] = None,
+    max_width: Optional[float] = None,
 ) -> str:
-    """Render a tensorgrad Tensor as book-style TikZ (uses tikz-styles.tex)."""
-    lines: list[str] = [
-        rf"\begin{{tikzpicture}}[baseline={baseline}, inner sep=1pt]"
-    ]
-    _emit_layout(layout_any(tensor, left, right), lines, prefix="", dx=0.0)
+    """Render a tensorgrad Tensor as book-style TikZ (uses tikz-styles.tex).
+
+    Args:
+        left/right: force the named free edge to exit that side (covariance).
+        baseline: TikZ baseline anchor for inline use.
+        scale: explicit TikZ scale factor for the whole picture.
+        max_width: if the laid-out diagram is wider than this (in cm), scale
+            it down to fit -- wide gradients/products then stay on the page
+            instead of overflowing. Ignored if `scale` is given.
+    """
+    layout = layout_any(tensor, left, right)
+    if scale is None and max_width is not None and layout.xmax > max_width > 0:
+        scale = max_width / layout.xmax
+    opts = f"baseline={baseline}, inner sep=1pt"
+    if scale is not None and abs(scale - 1.0) > 1e-6:
+        # `transform shape` scales node glyphs too, so the whole diagram
+        # shrinks uniformly instead of nodes overlapping at moved coordinates
+        opts += f", scale={scale:.3f}, transform shape"
+    lines: list[str] = [rf"\begin{{tikzpicture}}[{opts}]"]
+    _emit_layout(layout, lines, prefix="", dx=0.0)
     lines.append(r"\end{tikzpicture}")
     return "\n".join(lines)

tensorgrad/tensor.py

@@ -131,6 +131,18 @@ def edges(self) -> KeysView[str]:
         """Returns an _ordered_ set of edge names"""
         return self.shape.keys()
 
+    def to_book_tikz(self, **kwargs) -> str:
+        """Render this tensor as a book-style tensor diagram (TikZ).
+
+        See :func:`tensorgrad.extras.book_layout.to_book_tikz` for options
+        (``left``/``right`` to fix free-edge sides, ``max_width`` to fit wide
+        diagrams). Call ``.simplify()`` first if the expression still contains
+        derivatives you want evaluated.
+        """
+        from tensorgrad.extras.book_layout import to_book_tikz
+
+        return to_book_tikz(self, **kwargs)
+
     @property
     def shape(self) -> dict[str, Symbol]:
         if not hasattr(self, "_shape"):

tests/extras/test_book_layout.py

@@ -324,3 +324,28 @@ def test_expectation_nested_in_sum():
     assert len(lay.boxes) == 1
     ids = set(nd.id for nd in lay.nodes)
     assert all(a in ids for a in lay.boxes[0][0])
+
+
+def test_zero_tensor():
+    from tensorgrad.tensor import Zero
+
+    assert "0" in to_book_tikz(Zero(i=n))  # zero vector -> labeled 0 node
+    assert to_book_tikz(Zero(i=n, j=n)).count(r"\node") >= 1
+    assert "0" in to_book_tikz(Zero())  # scalar zero
+
+
+def test_max_width_scales_wide_diagrams():
+    import tensorgrad.functions as F
+
+    z = Variable("z", i=n)
+    y = Variable("y", i=n)
+    ce = F.cross_entropy(z, y, dim="i").grad(z, {"i": "j"}).simplify()
+    wide = to_book_tikz(ce)
+    narrowed = to_book_tikz(ce, max_width=8)
+    assert "transform shape" in narrowed and "scale=" in narrowed
+    assert "transform shape" not in wide  # unconstrained stays full size
+
+
+def test_scale_explicit():
+    tex = to_book_tikz(Product([_A(), Variable("B", j=n, k=n)]), scale=0.5)
+    assert "scale=0.5" in tex and "transform shape" in tex