model

Classes

AirModel

In-memory graph IR for AOT codegen.

Attributes:

• name (str) –

  Name of the model.

• version (str) –

  Version of the model.

• description (str) –

  Description of the model.

• tensors (dict[str, AirTensor]) –

  Mapping from tensor ID to its AirTensor metadata.

• operators (list[AirOperator]) –

  Topologically ordered list of operators (nodes) in the graph.

• input_ids (list[str]) –

  IDs of tensors that serve as model inputs.

• output_ids (list[str]) –

  IDs of tensors that serve as model outputs.

Attributes

inputs property

inputs: list[AirTensor]

List of AirTensor objects corresponding to model inputs.

outputs property

outputs: list[AirTensor]

List of AirTensor objects corresponding to model outputs.

Functions

add_tensor

add_tensor(tensor: AirTensor) -> None

Register a new tensor in the model.

Raises:

• ValueError –

  If a tensor with the same ID already exists.

remove_tensor

remove_tensor(tid: TensorId) -> None

Remove a tensor from the model by its ID.

Raises:

• KeyError –

  If the tensor ID is not found.

add_operator

add_operator(op: AirOperator) -> None

Append a new operator to the execution list.

Raises:

• ValueError –

  If an operator with the same ID already exists.

has_tensor

has_tensor(tid: TensorId) -> bool

Check if a tensor exists in the model by its ID.

get_tensor

get_tensor(tid: TensorId) -> AirTensor

Lookup a tensor by its ID.

Raises:

• KeyError –

  If the tensor ID is not found.

topo_sort

topo_sort() -> list[AirOperator]

Return operators in execution order.

Assumes operators was populated in valid topological order.

prune_unused_tensors

prune_unused_tensors() -> list[TensorId]

Remove dangling tensors that are not referenced by any operator or model input/output.