Skip to content

Visualization

scaly.viz records what the compiler does to chosen functions: the expression graph, the program after each optimization pass, and the generated C. It is opt-in. Importing scaly does not import it, and nothing is recorded for a function that was not marked.

import scaly as sc
from scaly.codegen import render_c_module
from scaly.viz import visualize

@sc.function(sc.arg("x", 3), outputs=sc.arg("energy"))
def energy(x: sc.Expr) -> sc.Expr:
    return sc.sumsqr(x)

visualize(energy, label="Energy model")
render_c_module(energy)

Every render of a marked function, whether ahead of time or through the JIT, adds one recording. Run uv run scaly_viz --browser to browse them. visualize is the same function as visualize_function.

Registering a declaration does not trace it. Functions with open shapes or no declarations record each instance when it renders, including instances created before registration. Registering a concrete instance records only that instance and takes precedence over an open declaration's registration. Removing either registration leaves the other in place. A fully specified declaration and its sole instance share one registration, so either can unregister it.

Recording

visualize_function 

visualize_function(fun: F, *, label: str | None = None) -> F

Mark fun for visualization on future AOT/JIT renders.

A declaration records all its current and future instances without tracing at registration. A concrete instance records only itself and takes precedence over an open declaration. Fully specified declarations share registration with their sole instance. The function itself is returned so callers can write f = visualize_function(f).

Source code in src/scaly/viz/recording.py 
49
50
51
52
53
54
55
56
57
58
59
def visualize_function[F: Function | ConcreteFunction](fun: F, *, label: str | None = None) -> F:
  """Mark ``fun`` for visualization on future AOT/JIT renders.

  A declaration records all its current and future instances without tracing at registration.
  A concrete instance records only itself and takes precedence over an open declaration.
  Fully specified declarations share registration with their sole instance.
  The function itself is returned so callers can write ``f = visualize_function(f)``.
  """
  with _RECORDING_LOCK:
    _VIZ_TARGETS[_target(fun)] = label
  return fun

unvisualize_function 

unvisualize_function(fun: Function | ConcreteFunction) -> None

Remove fun's registration. Other declarations or instances remain registered.

A fully specified declaration and its sole instance share one registration.

Source code in src/scaly/viz/recording.py 
62
63
64
65
66
67
68
def unvisualize_function(fun: Function | ConcreteFunction) -> None:
  """Remove ``fun``'s registration. Other declarations or instances remain registered.

  A fully specified declaration and its sole instance share one registration.
  """
  with _RECORDING_LOCK:
    _VIZ_TARGETS.pop(_target(fun), None)

capture 

capture(fun: F, *, label: str | None = None) -> Generator[F]

Record renders of fun only inside a with block, which yields fun.

Source code in src/scaly/viz/recording.py 
71
72
73
74
75
76
77
78
@contextmanager
def capture[F: Function | ConcreteFunction](fun: F, *, label: str | None = None) -> Generator[F]:
  """Record renders of ``fun`` only inside a ``with`` block, which yields ``fun``."""
  visualize_function(fun, label=label)
  try:
    yield fun
  finally:
    unvisualize_function(fun)

Reading recordings

recordings 

recordings() -> list[dict[str, Any]]

Return the recordings made in this process, oldest first.

Each is a JSON-ready dict with the display name, the function name, any render error and the steps: the expression graph, each program-dialect pass and the generated C.

Source code in src/scaly/viz/recording.py 
89
90
91
92
93
94
95
96
def recordings() -> list[dict[str, Any]]:
  """Return the recordings made in this process, oldest first.

  Each is a JSON-ready dict with the display ``name``, the ``function`` name, any render ``error``
  and the ``steps``: the expression graph, each program-dialect pass and the generated C.
  """
  with _RECORDING_LOCK:
    return list(_RECORDINGS)

load_recordings 

load_recordings(path: str | PathLike[str] | None = None) -> list[dict[str, Any]]

Read the recordings saved at path, by default recording_path(), including those from other processes.

Source code in src/scaly/viz/recording.py 
 99
100
101
102
103
104
def load_recordings(path: str | os.PathLike[str] | None = None) -> list[dict[str, Any]]:
  """Read the recordings saved at ``path``, by default ``recording_path()``, including those from other processes."""
  p = Path(path) if path is not None else recording_path()
  if not p.exists():
    return []
  return json.loads(p.read_text())

clear_recordings 

clear_recordings(*, disk: bool = False) -> None

Forget the recordings made in this process, and with disk=True delete the recording file too.

Source code in src/scaly/viz/recording.py 
81
82
83
84
85
86
def clear_recordings(*, disk: bool = False) -> None:
  """Forget the recordings made in this process, and with ``disk=True`` delete the recording file too."""
  with _RECORDING_LOCK:
    _RECORDINGS.clear()
  if disk:
    recording_path().unlink(missing_ok=True)

recording_dir 

recording_dir() -> Path

Return the directory recordings are written to: SCALY_VIZ_DIR, else $XDG_CACHE_HOME/scaly/viz, else ~/.cache/scaly/viz.

Source code in src/scaly/viz/recording.py 
28
29
30
31
32
33
34
35
def recording_dir() -> Path:
  """Return the directory recordings are written to: ``SCALY_VIZ_DIR``, else ``$XDG_CACHE_HOME/scaly/viz``, else ``~/.cache/scaly/viz``."""
  override = os.environ.get("SCALY_VIZ_DIR")
  if override:
    return Path(override).expanduser()
  xdg = os.environ.get("XDG_CACHE_HOME")
  base = Path(xdg).expanduser() if xdg else Path.home() / ".cache"
  return base / "scaly" / "viz"

recording_path 

recording_path() -> Path

Return the JSON file that every finished recording is appended to.

Source code in src/scaly/viz/recording.py 
38
39
40
def recording_path() -> Path:
  """Return the JSON file that every finished recording is appended to."""
  return recording_dir() / "recordings.json"

Browsing

serve 

serve(*, host: str = '127.0.0.1', port: int = 8000, path: str | Path | None = None, open_browser: bool = False) -> None

Serve the recordings at path, by default recording_path(), until interrupted. The scaly_viz command runs this.

Source code in src/scaly/viz/serve.py 
384
385
386
387
388
389
390
391
392
393
394
395
def serve(*, host: str = "127.0.0.1", port: int = 8000, path: str | Path | None = None, open_browser: bool = False) -> None:
  """Serve the recordings at ``path``, by default ``recording_path()``, until interrupted. The ``scaly_viz`` command runs this."""
  recording_file = Path(path) if path is not None else recording_path()
  server = VizServer((host, port), recording_file)
  url = f"http://{host}:{port}"
  print(f"*** scaly viz serving {recording_file} on {url}")
  if open_browser:
    threading.Timer(0.2, lambda: webbrowser.open(url)).start()
  try:
    server.serve_forever()
  except KeyboardInterrupt:
    print("*** scaly viz shutting down")