Weights¶

tinax.weights owns tensor manifests and bounded Safetensors interchange. Loading inspects manifests and byte budgets before tensor materialization. Serialization accepts explicit host tensors and never silently gathers global JAX arrays.

See the Safetensors Interchange guide for task-oriented usage.

TensorManifest `dataclass` ¶

TensorManifest(rules: tuple[TensorRule, ...])

Apply immutable rules with exact source coverage.

Attributes:

Name	Type	Description
`rules`	`tuple[TensorRule, ...]`	Tuple of `TensorRule` with distinct source and destination names.

Raises:

Type	Description
`TypeError`	If `rules` is not a tuple of `TensorRule` instances.
`ValueError`	If two rules share a source name or collide on a destination name.

apply ¶

apply(
    tensors: Mapping[str, NDArray[generic]],
) -> dict[str, NDArray[generic]]

Transform tensors after requiring exact source-name coverage.

Parameters:

Name	Type	Description	Default
`tensors`	`Mapping[str, NDArray[generic]]`	Mapping of source name to host NumPy array. Its key set must exactly match the manifest's source names.	required

Returns:

Type	Description
`dict[str, NDArray[generic]]`	A dict mapping each destination name to its transformed NumPy array.

Raises:

Type	Description
`TypeError`	If `tensors` is not a mapping, a source is not a NumPy array, or a transform returns the wrong type or dtype.
`ValueError`	If the source names do not exactly match, or a transform produces an unexpected shape.

TensorRule `dataclass` ¶

TensorRule(
    source_name: str,
    destination_name: str,
    transform: Callable[
        [NDArray[generic]], NDArray[generic]
    ],
    expected_shape: tuple[int, ...],
    expected_dtype: dtype[generic],
)

Describe one source-to-destination tensor transformation.

Attributes:

Name	Type	Description
`source_name`	`str`	Non-empty source tensor name.
`destination_name`	`str`	Non-empty destination tensor name.
`transform`	`Callable[[NDArray[generic]], NDArray[generic]]`	Callable mapping the source NumPy array to the destination array.
`expected_shape`	`tuple[int, ...]`	Tuple of non-negative dimensions the transform must produce.
`expected_dtype`	`dtype[generic]`	NumPy dtype the transform must produce.

Raises:

Type	Description
`TypeError`	If a name is not a string, `transform` is not callable, `expected_shape` is not a tuple of integers, or `expected_dtype` is not a valid NumPy dtype.
`ValueError`	If a name is empty or an `expected_shape` dimension is negative.

inspect_safetensors ¶

inspect_safetensors(
    path: str | PathLike[str],
    *,
    max_metadata_bytes: int = 0,
) -> SafetensorsInfo

Inspect tensor keys, shapes, dtypes, sizes, and bounded metadata without loading payloads.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Path to the Safetensors file.	required
`max_metadata_bytes`	`int`	Maximum header metadata bytes to accept. `0` allows none.	`0`

Returns:

Type	Description
`SafetensorsInfo`	A `SafetensorsInfo` describing the file's tensors and validated metadata.

Raises:

Type	Description
`TypeError`	If `max_metadata_bytes` is not an integer.
`ValueError`	If `max_metadata_bytes` is negative, the metadata exceeds the budget, or a tensor name is invalid or duplicated.

load_safetensors ¶

load_safetensors(
    path: str | PathLike[str],
    *,
    max_bytes: int,
    names: Iterable[str] | None = None,
    max_metadata_bytes: int = 0,
) -> LoadedSafetensors

Inspect and selectively load NumPy tensors within an explicit nominal byte budget.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Path to the Safetensors file.	required
`max_bytes`	`int`	Maximum total nominal payload bytes to load.	required
`names`	`Iterable[str] \| None`	Tensor names to load, or `None` to load all present tensors.	`None`
`max_metadata_bytes`	`int`	Maximum header metadata bytes to accept. `0` allows none.	`0`

Returns:

Type	Description
`LoadedSafetensors`	A `LoadedSafetensors` holding the selected host tensors and file information.

Raises:

Type	Description
`TypeError`	If a byte limit is not an integer, or a loaded tensor has an unmaterializable or unexpected dtype.
`ValueError`	If a byte limit is negative, requested names are absent, the selection exceeds `max_bytes`, or a loaded tensor has an unexpected shape.

save_safetensors ¶

save_safetensors(
    path: str | PathLike[str],
    tensors: Mapping[str, NDArray[generic]],
    *,
    overwrite: bool,
    metadata: Mapping[str, str] | None = None,
    max_metadata_bytes: int = 0,
) -> None

Atomically save explicit C-contiguous host tensors in the destination directory.

Parameters:

Name	Type	Description	Default
`path`	`str \| PathLike[str]`	Destination file path. Written atomically through a temporary file.	required
`tensors`	`Mapping[str, NDArray[generic]]`	Mapping of tensor name to host NumPy array to serialize.	required
`overwrite`	`bool`	Whether to replace an existing destination.	required
`metadata`	`Mapping[str, str] \| None`	Optional string-to-string header metadata.	`None`
`max_metadata_bytes`	`int`	Maximum metadata bytes to accept. `0` allows none.	`0`

Raises:

Type	Description
`TypeError`	If `overwrite` is not a bool, `max_metadata_bytes` is not an integer, or `tensors` or `metadata` have invalid types.
`ValueError`	If `max_metadata_bytes` is negative, the metadata exceeds the budget, a tensor is invalid, or the destination exists and `overwrite` is `False`.

SafetensorsInfo `dataclass` ¶

SafetensorsInfo(
    tensors: Mapping[str, TensorInfo],
    metadata: Mapping[str, str],
)

Describe a validated Safetensors header and its bounded metadata.

Attributes:

Name	Type	Description
`tensors`	`Mapping[str, TensorInfo]`	Read-only mapping of tensor name to `TensorInfo`.
`metadata`	`Mapping[str, str]`	Read-only mapping of validated header metadata.

Raises:

Type	Description
`TypeError`	If `tensors` is not a mapping or a value is not a `TensorInfo`.
`ValueError`	If a tensor name is invalid.

total_bytes `property` ¶

total_bytes: int

Return the nominal payload size of all tensors.

TensorInfo `dataclass` ¶

TensorInfo(shape: tuple[int, ...], dtype: str)

Describe one tensor without materializing its payload.

Attributes:

Name	Type	Description
`shape`	`tuple[int, ...]`	Tuple of non-negative tensor dimensions.
`dtype`	`str`	Safetensors dtype string (for example, `"F32"`).

Raises:

Type	Description
`TypeError`	If `shape` is not a tuple of integers or `dtype` is not a string.
`ValueError`	If a `shape` dimension is negative or the dtype and shape are invalid.

nbytes `property` ¶

nbytes: int

Return the nominal serialized payload size.

LoadedSafetensors `dataclass` ¶

LoadedSafetensors(
    tensors: Mapping[str, NDArray[generic]],
    info: SafetensorsInfo,
)

Hold selected host tensors and the inspected file information.

Attributes:

Name	Type	Description
`tensors`	`Mapping[str, NDArray[generic]]`	Read-only mapping of tensor name to host NumPy array.
`info`	`SafetensorsInfo`	The `SafetensorsInfo` the tensors were validated against.

Raises:

Type	Description
`TypeError`	If `info` is not a `SafetensorsInfo`, `tensors` is not a mapping, a tensor is not a NumPy array, or a dtype disagrees with `info`.
`ValueError`	If a tensor name is invalid, absent from `info`, or has an unexpected shape.

metadata `property` ¶

metadata: Mapping[str, str]

Return the validated metadata from the file header.

Weights¶

TensorManifest dataclass ¶

apply ¶

TensorRule dataclass ¶

inspect_safetensors ¶

load_safetensors ¶

save_safetensors ¶

SafetensorsInfo dataclass ¶

total_bytes property ¶

TensorInfo dataclass ¶

nbytes property ¶

LoadedSafetensors dataclass ¶

metadata property ¶

TensorManifest `dataclass` ¶

TensorRule `dataclass` ¶

SafetensorsInfo `dataclass` ¶

total_bytes `property` ¶

TensorInfo `dataclass` ¶

nbytes `property` ¶

LoadedSafetensors `dataclass` ¶

metadata `property` ¶