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.L("x", 3), sc.L("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.

Recording¶

visualize_function ¶

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

Mark fun for visualization on future AOT/JIT renders.

Capturing is exact-object opt-in: no Function is recorded unless it was passed here (or to capture). The function itself is returned so callers can write f = visualize_function(f).

Source code in src/scaly/viz/recording.py

def visualize_function(fun: Function, *, label: str | None = None) -> Function:
  """Mark ``fun`` for visualization on future AOT/JIT renders.

  Capturing is exact-object opt-in: no Function is recorded unless it was passed
  here (or to ``capture``). The function itself is returned so callers can write
  ``f = visualize_function(f)``.
  """
  with _RECORDING_LOCK:
    _VIZ_TARGETS[id(fun)] = label
  return fun

unvisualize_function ¶

unvisualize_function(fun: Function) -> None

Stop recording renders of fun. Does nothing if it was not marked.

Source code in src/scaly/viz/recording.py

def unvisualize_function(fun: Function) -> None:
  """Stop recording renders of ``fun``. Does nothing if it was not marked."""
  with _RECORDING_LOCK:
    _VIZ_TARGETS.pop(id(fun), None)

capture ¶

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

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

Source code in src/scaly/viz/recording.py

@contextmanager
def capture(fun: Function, *, label: str | None = None) -> Iterator[Function]:
  """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

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

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

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

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

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

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")