ModelTensorDefinition

SUMMARY

The definition of a model input or output tensor.

from telekinesis import datatypes
tensor = datatypes.ModelTensorDefinition(
    name="images", canonical_name="input_0", dtype="float32", shape=[1, 3, 224, 224]
)

Not directly serializable

Although registered as a BaseDataType, its to_pyarrow/from_pyarrow are batch-oriented classmethods (a list[ModelTensorDefinition] in, one StructArray out) rather than the standard per-instance contract. ModelTensorDefinition is not meant to be passed directly to datatypes.serialize(). It's serialized only as an element of ModelDefinitions.model_inputs/model_outputs.

API Reference

Complete API documentation for ModelTensorDefinition, including parameters, attributes, and methods.

View Reference →

Parameters

Parameter	Type	Default	Description
name	str	Required	Tensor name used by the model.
canonical_name	str	Required	Canonical name used to identify the tensor consistently across model formats or runtimes.
dtype	str	Required	Tensor element dtype, such as float32.
shape	list[int] | tuple[int, ...]	Required	Tensor dimensions. Use -1 for a dynamic dimension.

Raises

Exception	Condition
TypeError	name/canonical_name/dtype is not a str, shape is not a list/tuple, or a shape entry is not an int.

Attributes

Attribute	Type	Description
name	str	Tensor name.
canonical_name	str	Canonical tensor name.
dtype	str	Tensor element dtype.
shape	list[int]	Tensor shape.

Reading shape also returns the stored list directly rather than a defensive copy, so mutating that list in place affects the instance.

Methods

Method	Type	Description
ModelTensorDefinition.coerce(value, name="value")	ModelTensorDefinition	Converts a dict into a ModelTensorDefinition via from_dict. If value is already a ModelTensorDefinition, it is returned unchanged.
ModelTensorDefinition.from_dict(value)	ModelTensorDefinition	Builds an instance from a dict containing the keys name, canonical_name, dtype, and shape; extra keys are ignored. All four keys are required, and each value must have the right type.
ModelTensorDefinition.from_pyarrow(col)	list[ModelTensorDefinition]	Inverse of to_pyarrow: deserializes a pa.StructArray back into a list of ModelTensorDefinition, preserving order. col's fields must match the schema returned by arrow_type().
to_dict()	dict	Returns the tensor definition as a plain dict with keys name, canonical_name, dtype, and shape. The shape value is a fresh list, so mutating it doesn't affect the instance.
ModelTensorDefinition.to_pyarrow(values)	pa.StructArray	Serializes a list of ModelTensorDefinition (which may be empty) into a single pa.StructArray matching the schema from arrow_type().
ModelTensorDefinition.arrow_type()	pa.StructType	Returns the struct type (fields name, canonical_name, dtype, and shape: list<int32>) used to serialize a batch to and from Arrow.

Operators

Operation	Behavior
a == b	True if other is a ModelTensorDefinition with equal name, canonical_name, dtype, and shape; NotImplemented for any other type.

Example

ModelTensorDefinition has no standalone example file. It's demonstrated below as the model_inputs/model_outputs entries of ModelDefinitions:

"""Demonstrates the Telekinesis ModelDefinitions datatype."""

import time
from datetime import datetime, timezone

import rerun as rr
from loguru import logger

from telekinesis import datatypes

def model_definitions_example():
    """Demonstrate creation, inspection, operations, visualization, and serialization."""

    # ======================= Create ============================================
    created_at = datetime(2024, 6, 1, tzinfo=timezone.utc)
    updated_at = datetime(2024, 6, 15, tzinfo=timezone.utc)

    model_input = datatypes.ModelTensorDefinition(
        name="images", canonical_name="input_0", dtype="float32", shape=[1, 3, 224, 224]
    )
    model_output = datatypes.ModelTensorDefinition(
        name="logits", canonical_name="output_0", dtype="float32", shape=[1, 1000]
    )
    other_output = datatypes.ModelTensorDefinition(
        name="logits", canonical_name="output_0", dtype="float32", shape=[1, 1000]
    )

    definitions = datatypes.ModelDefinitions(
        model_names=["model-a", "model-b"],
        model_formats=["onnx", "pytorch"],
        visibilities=["private", "public"],
        model_statuses=["uploaded", "deploying"],
        model_descriptions=["first model", None],
        model_inputs=[[model_input], None],
        model_outputs=[[model_output, other_output], None],
        created_ats=[created_at, updated_at],
        updated_ats=[created_at, updated_at],
    )
    logger.info(f"Created ModelDefinitions: {definitions}")

    # ======================= Inspect ===========================================
    logger.info(f"model_names={definitions.model_names}")
    logger.info(f"model_formats={definitions.model_formats}")
    logger.info(f"visibilities={definitions.visibilities}")
    logger.info(f"model_statuses={definitions.model_statuses}")
    logger.info(f"model_descriptions={definitions.model_descriptions}")
    logger.info(f"model_inputs={definitions.model_inputs}")
    logger.info(f"model_outputs={definitions.model_outputs}")
    logger.info(f"created_ats={definitions.created_ats}")
    logger.info(f"updated_ats={definitions.updated_ats}")

    # ======================= Operations =========================================
    logger.info(f"length={len(definitions)}")

    subset = definitions[0:1]
    logger.info(f"definitions[0:1] = {len(subset)} record(s), names={subset.model_names}")

    first = definitions[0]
    logger.info(f"definitions[0] = {first}")

    mask = definitions.model_statuses == "uploaded"
    logger.info(f"definitions[uploaded mask] = {len(definitions[mask])} record(s)")

    # ModelTensorDefinition exposes its fields as plain attributes plus a dict view.
    logger.info(f"model_input.to_dict()={model_input.to_dict()}")

    # A ModelDefinitions batch may be empty.
    empty = datatypes.ModelDefinitions(
        model_names=[], model_formats=[], visibilities=[], model_statuses=[]
    )
    logger.info(f"length of empty batch={len(empty)}")

    # ======================= Visualize =========================================
    rr.init("model_definitions_example", spawn=True)
    datatypes.visualize(definitions, entity_path="/model_definitions")
    datatypes.visualize(empty, entity_path="/model_definitions/empty")

    # ======================= Serialize / Deserialize ===========================
    start = time.perf_counter()
    serialized = datatypes.serialize(definitions)
    serialization_ms = (time.perf_counter() - start) * 1000

    start = time.perf_counter()
    deserialized = datatypes.deserialize(serialized)["param_0"]
    deserialization_ms = (time.perf_counter() - start) * 1000

    logger.info(f"Deserialized ModelDefinitions: {deserialized}")
    logger.info(f"Round-trip successful: {definitions == deserialized}")
    logger.info(f"Serialization time: {serialization_ms:.3f} ms")
    logger.info(f"Deserialization time: {deserialization_ms:.3f} ms")


if __name__ == "__main__":
    model_definitions_example()