trama 0.2.0

JSON format for composable model graphs

Trama

Trama is a standalone Python library that defines a JSON schema, validator, and export/import framework for hierarchical model graphs in Static Single Assignment (SSA) style. It is framework-agnostic and is designed to be used by domain-specific modeling libraries that inject their own op registries. Models are serialized to self-contained .trama.json files capturing both topology and parameter values, and can be reconstructed back into runtime modules via registered factory functions.

Installation

# With pip:
pip install trama

# Or with uv:
uv add trama

Quick-start

import trama

# Export a module that implements ExportableOp + ImportableOp.
model = trama.export_model(module, inputs=["x"], outputs=["y"], registry=registry)

# Serialize to JSON and save.
trama.dump(model, "my_model.trama.json")

# Later: load and reconstruct the runtime module.
model = trama.load("my_model.trama.json")
instance = trama.import_model(model, registry=registry)

JSON Schema

Trama exposes a JSON Schema for .trama.json files. Generate it with the tramaschema command:

tramaschema > trama.schema.json

The schema validates format structure: required fields, field types, and identifier patterns. Registry conformance — whether op names resolve and parameter types match — is not expressible in JSON Schema and is enforced at load time in Python by trama.load() or trama.validate().

Example

{
  "format_version": "0.2",
  "inputs": ["x"],
  "outputs": [{ "node": "root", "output": "y" }],
  "nodes": [
    {
      "id": "root", "parent": null, "op": "Sequential",
      "inputs": { "x": { "node": "%input", "output": "x" } },
      "outputs": ["y"],
      "parameters": {},
      "output_bindings": { "y": { "node": "root.linear_1", "output": "y" } }
    },
    {
      "id": "root.linear_0", "parent": "root", "op": "Linear",
      "inputs": { "x": { "node": "%input", "output": "x" } },
      "outputs": ["y"],
      "parameters": {
        "in_features":  { "name": "root.linear_0.in_features" },
        "out_features": { "name": "root.linear_0.out_features" },
        "weight":       { "name": "root.linear_0.weight" }
      },
      "output_bindings": null
    },
    {
      "id": "root.linear_1", "parent": "root", "op": "Linear",
      "inputs": { "x": { "node": "root.linear_0", "output": "y" } },
      "outputs": ["y"],
      "parameters": {
        "in_features":  { "name": "root.linear_0.out_features" },
        "out_features": { "name": "root.linear_1.out_features" },
        "weight":       { "name": "root.linear_1.weight" }
      },
      "output_bindings": null
    }
  ],
  "parameters": {
    "root.linear_0.in_features":  { "value": 8,  "learnable": false },
    "root.linear_0.out_features": { "value": 16, "learnable": false },
    "root.linear_1.out_features": { "value": 4,  "learnable": false },
    "root.linear_0.weight": { "value": { "dtype": "float32", "shape": [16, 8],  "data": "..." }, "learnable": true },
    "root.linear_1.weight": { "value": { "dtype": "float32", "shape": [4,  16], "data": "..." }, "learnable": true }
  }
}

Note that root.linear_0.out_features is referenced by both root.linear_0 and root.linear_1 (the dimension is the same physical value). This is the parameter sharing mechanism: a single entry in the table, multiple refs.

The trama format is inspired by SSA form from compiler design. The tramassa command prints any .trama.json file in this human-readable form:

# trama SSA  (format 0.2)

model(x) -> %root.y

params
  root.linear_0.in_features   : int         = 8
  root.linear_0.out_features  : int         = 16
  root.linear_1.out_features  : int         = 4
  root.linear_0.weight        : f32[16, 8]  learnable
  root.linear_1.weight        : f32[4, 16]  learnable

scope root [Sequential]
  y = %root.linear_1.y

%root.linear_0.y = Linear(
    x            = x,
    in_features  = root.linear_0.in_features,
    out_features = root.linear_0.out_features,
    weight       = root.linear_0.weight
)

%root.linear_1.y = Linear(
    x            = %root.linear_0.y,
    in_features  = root.linear_0.out_features,
    out_features = root.linear_1.out_features,
    weight       = root.linear_1.weight
)

Design

For detailed description of the design of Trama please see the design document.

Compared to other model export formats

Trama operates at a fundamentally different abstraction level than ONNX, torch.export, and similar formats. Those tools target deployment: they decompose user-authored modules into low-level tensor operations (aten.matmul, aten.conv2d, ONNX Gemm) so a runtime can execute them efficiently. The user's Sequential(Encoder(), Decoder()) becomes a flat sequence of primitive ops with the original module boundaries dissolved. Trama instead targets composition: it preserves the user-authored module hierarchy as first-class structure, captures only the wiring between the building blocks the consumer library defines, and treats those blocks as opaque ops whose semantics live in the registry. This makes trama files dramatically smaller and more readable than ONNX graphs of equivalent models, diffable in code review, and meaningful to humans inspecting "what does this model architecture actually look like" — at the cost of being unsuitable for cross-framework deployment, since the ops are defined by the consumer library rather than a standardized vocabulary. Trama is complementary to ONNX, not a replacement: a natural pipeline is to author and review models in trama, then lower to ONNX or torch.export when shipping to inference.