tensor

Adapt Arx surface tensor syntax to IRx Tensor nodes while keeping user-facing shape and indexing rules local to Arx.

Classes:

• TensorBinding –

Functions:

• attach_binding –
• binding_from_type –
• build_literal_from_literal –
• coerce_expression –
• default_value –
• infer_literal –
• is_tensor_type –
• literal_values –
• runtime_tensor_type –
• tensor_shape –
• tensor_type –

TensorBinding dataclass

TensorBinding(
    element_type: DataType, layout: TensorLayout, flags: int
)

attach_binding

attach_binding(node: AST, binding: TensorBinding) -> None

Source code in packages/arx/src/arx/tensor.py

def attach_binding(node: astx.AST, binding: TensorBinding) -> None:
    """
    title: Attach static tensor metadata to one AST node.
    parameters:
      node:
        type: astx.AST
      binding:
        type: TensorBinding
    """
    info = cast(SemanticInfo | None, getattr(node, "semantic", None))
    if info is None or not isinstance(info, SemanticInfo):
        info = SemanticInfo()
        setattr(node, "semantic", info)
    info.extras[TENSOR_LAYOUT_EXTRA] = binding.layout
    info.extras[TENSOR_ELEMENT_TYPE_EXTRA] = binding.element_type
    info.extras[TENSOR_FLAGS_EXTRA] = binding.flags

binding_from_type

binding_from_type(
    data_type: DataType | None,
) -> TensorBinding | None

Source code in packages/arx/src/arx/tensor.py

def binding_from_type(
    data_type: astx.DataType | None,
) -> TensorBinding | None:
    """
    title: Build one static tensor binding from one declared type.
    parameters:
      data_type:
        type: astx.DataType | None
    returns:
      type: TensorBinding | None
    """
    if not is_tensor_type(data_type):
        return None

    shape = tensor_shape(data_type)
    element_type = cast(astx.TensorType, data_type).element_type
    if shape is None or element_type is None:
        return None

    item_size = _element_size_bytes(element_type)
    layout = TensorLayout(
        shape=shape,
        strides=tensor_default_strides(shape, item_size),
        offset_bytes=0,
    )
    flags = buffer_view_flags(
        BufferOwnership.EXTERNAL_OWNER,
        BufferMutability.READONLY,
        c_contiguous=tensor_is_c_contiguous(layout, item_size),
        f_contiguous=tensor_is_f_contiguous(layout, item_size),
    )
    return TensorBinding(
        element_type=element_type,
        layout=layout,
        flags=flags,
    )

build_literal_from_literal

build_literal_from_literal(
    expr: Literal, target_type: DataType, *, context: str
) -> TensorLiteral

Source code in packages/arx/src/arx/tensor.py

def build_literal_from_literal(
    expr: astx.Literal,
    target_type: astx.DataType,
    *,
    context: str,
) -> astx.TensorLiteral:
    """
    title: Build one tensor literal from one nested literal value.
    parameters:
      expr:
        type: astx.Literal
      target_type:
        type: astx.DataType
      context:
        type: str
    returns:
      type: astx.TensorLiteral
    """
    binding = binding_from_type(target_type)
    if binding is None:
        if not is_tensor_type(target_type):
            raise ValueError("tensor literal target must be a tensor type")
        raise ValueError(
            "tensor literal target requires a static tensor shape"
        )

    shape, values = _flatten_literal(expr)
    if shape != binding.layout.shape:
        raise ValueError(
            f"{context} has shape {_format_shape(shape)} but the declared "
            f"tensor shape is {_format_shape(binding.layout.shape)}"
        )

    for value in values:
        _validate_scalar_literal(value, binding.element_type, context=context)

    return _literal_from_values(binding, values)

coerce_expression

coerce_expression(
    expr: Expr, target_type: DataType, *, context: str
) -> Expr

Source code in packages/arx/src/arx/tensor.py

def coerce_expression(
    expr: astx.Expr,
    target_type: astx.DataType,
    *,
    context: str,
) -> astx.Expr:
    """
    title: Coerce one parsed expression into one declared tensor type.
    parameters:
      expr:
        type: astx.Expr
      target_type:
        type: astx.DataType
      context:
        type: str
    returns:
      type: astx.Expr
    """
    if not is_tensor_type(target_type):
        return expr
    if isinstance(expr, astx.TensorLiteral):
        return expr
    if not isinstance(expr, astx.Literal):
        return expr
    return build_literal_from_literal(expr, target_type, context=context)

default_value

default_value(target_type: DataType) -> TensorLiteral

Source code in packages/arx/src/arx/tensor.py

def default_value(target_type: astx.DataType) -> astx.TensorLiteral:
    """
    title: Build one default tensor literal for one declared type.
    parameters:
      target_type:
        type: astx.DataType
    returns:
      type: astx.TensorLiteral
    """
    binding = binding_from_type(target_type)
    if binding is None:
        raise ValueError("default tensor value requires a static tensor shape")

    count = prod(binding.layout.shape)
    scalar = _zero_literal(binding.element_type)
    values = tuple(_clone_scalar_literal(scalar) for _ in range(count))
    return _literal_from_values(binding, values)

infer_literal

infer_literal(expr: Literal) -> TensorLiteral

Source code in packages/arx/src/arx/tensor.py

def infer_literal(expr: astx.Literal) -> astx.TensorLiteral:
    """
    title: Infer one tensor literal directly from one literal value.
    parameters:
      expr:
        type: astx.Literal
    returns:
      type: astx.TensorLiteral
    """
    shape, values = _flatten_literal(expr)
    if not values:
        raise ValueError(
            "cannot infer a tensor element type from an empty literal"
        )

    element_type = _infer_element_type(values[0])
    for value in values:
        _validate_scalar_literal(
            value,
            element_type,
            context="tensor literal",
        )

    binding = cast(
        TensorBinding,
        binding_from_type(tensor_type(element_type, shape)),
    )
    return _literal_from_values(binding, values)

is_tensor_type

is_tensor_type(data_type: DataType | None) -> bool

Source code in packages/arx/src/arx/tensor.py

def is_tensor_type(data_type: astx.DataType | None) -> bool:
    """
    title: Return whether one type is an Arx tensor surface type.
    parameters:
      data_type:
        type: astx.DataType | None
    returns:
      type: bool
    """
    return (
        isinstance(data_type, astx.TensorType)
        and getattr(data_type, TENSOR_SURFACE_ATTR, False) is True
    )

literal_values

literal_values(node: TensorLiteral) -> tuple[AST, ...]

Source code in packages/arx/src/arx/tensor.py

def literal_values(
    node: astx.TensorLiteral,
) -> tuple[astx.AST, ...]:
    """
    title: Return one flattened scalar payload from a tensor literal.
    parameters:
      node:
        type: astx.TensorLiteral
    returns:
      type: tuple[astx.AST, Ellipsis]
    """
    return tuple(node.values)

runtime_tensor_type

runtime_tensor_type(element_type: DataType) -> TensorType

Source code in packages/arx/src/arx/tensor.py

def runtime_tensor_type(element_type: astx.DataType) -> astx.TensorType:
    """
    title: Build one runtime-shaped tensor surface type.
    parameters:
      element_type:
        type: astx.DataType
    returns:
      type: astx.TensorType
    """
    _element_size_bytes(element_type)
    return _mark_tensor_type(astx.TensorType(element_type), None)

tensor_shape

tensor_shape(
    data_type: DataType | None,
) -> tuple[int, ...] | None

Source code in packages/arx/src/arx/tensor.py

def tensor_shape(data_type: astx.DataType | None) -> tuple[int, ...] | None:
    """
    title: Return the declared tensor shape when available.
    parameters:
      data_type:
        type: astx.DataType | None
    returns:
      type: tuple[int, Ellipsis] | None
    """
    if data_type is None:
        return None
    return _shape_of(data_type)

tensor_type

tensor_type(
    element_type: DataType, shape: tuple[int, ...]
) -> TensorType

Source code in packages/arx/src/arx/tensor.py

def tensor_type(
    element_type: astx.DataType,
    shape: tuple[int, ...],
) -> astx.TensorType:
    """
    title: Build one static-shape tensor surface type.
    parameters:
      element_type:
        type: astx.DataType
      shape:
        type: tuple[int, Ellipsis]
    returns:
      type: astx.TensorType
    """
    _element_size_bytes(element_type)
    if not shape:
        raise ValueError("tensor shapes must include at least one dimension")
    if any(dim < 0 for dim in shape):
        raise ValueError("tensor dimensions must be non-negative")
    return _mark_tensor_type(astx.TensorType(element_type), shape)