"""
Persistence utilities for SparseTensor.
* ``safetensors`` format for efficient, safe serialisation.
* Matrix Market (``.mtx``) for interop with SciPy and other tools.
::
from torch_sla import SparseTensor
from torch_sla.io import save_sparse, load_sparse
A = SparseTensor(val, row, col, shape)
save_sparse(A, "matrix.safetensors")
A = load_sparse("matrix.safetensors")
"""
import os
import json
import torch
from typing import Dict, Optional, Tuple, Union, TYPE_CHECKING
from pathlib import Path
if TYPE_CHECKING:
from .sparse_tensor import SparseTensor
try:
from safetensors.torch import save_file, load_file
SAFETENSORS_AVAILABLE = True
except ImportError:
SAFETENSORS_AVAILABLE = False
def _ensure_safetensors():
"""Ensure safetensors is available."""
if not SAFETENSORS_AVAILABLE:
raise ImportError(
"safetensors is required for persistence. "
"Install with: pip install safetensors"
)
# =============================================================================
# Matrix Market (.mtx) Format
# =============================================================================
[docs]
def save_mtx(
tensor: "SparseTensor",
path: Union[str, Path],
comment: str = "",
field: str = "real",
symmetry: str = "general",
) -> None:
"""
Save a SparseTensor to Matrix Market (.mtx) format.
Parameters
----------
tensor : SparseTensor
The sparse tensor to save.
path : str or Path
Output file path (should end with .mtx).
comment : str, optional
Comment to include in the header.
field : str, optional
Field type: 'real', 'complex', 'integer', or 'pattern'.
Default: 'real'.
symmetry : str, optional
Symmetry type: 'general', 'symmetric', 'skew-symmetric', or 'hermitian'.
Default: 'general'.
Example
-------
>>> A = SparseTensor(val, row, col, (100, 100))
>>> save_mtx(A, "matrix.mtx")
>>> save_mtx(A, "matrix.mtx", symmetry="symmetric")
"""
from .sparse_tensor import SparseTensor
if not isinstance(tensor, SparseTensor):
raise TypeError(f"Expected SparseTensor, got {type(tensor)}")
if tensor.is_batched:
raise ValueError("Cannot save batched SparseTensor to Matrix Market format")
path = Path(path)
M, N = tensor.sparse_shape
nnz = tensor.nnz
# Get data on CPU
row = tensor.row_indices.cpu().numpy()
col = tensor.col_indices.cpu().numpy()
val = tensor.values.cpu().numpy()
# Determine field type from dtype
if tensor.dtype in (torch.complex64, torch.complex128):
field = "complex"
elif tensor.dtype in (torch.int32, torch.int64):
field = "integer"
with open(path, 'w') as f:
# Write header
f.write(f"%%MatrixMarket matrix coordinate {field} {symmetry}\n")
if comment:
for line in comment.split('\n'):
f.write(f"% {line}\n")
f.write(f"% Generated by torch-sla\n")
# Write dimensions
f.write(f"{M} {N} {nnz}\n")
# Write entries (1-indexed)
if field == "pattern":
for i in range(nnz):
f.write(f"{row[i] + 1} {col[i] + 1}\n")
elif field == "complex":
for i in range(nnz):
v = val[i]
f.write(f"{row[i] + 1} {col[i] + 1} {v.real} {v.imag}\n")
else:
for i in range(nnz):
f.write(f"{row[i] + 1} {col[i] + 1} {val[i]}\n")
[docs]
def load_mtx(
path: Union[str, Path],
dtype: Optional[torch.dtype] = None,
device: Union[str, torch.device] = "cpu",
) -> "SparseTensor":
"""
Load a SparseTensor from Matrix Market (.mtx) format.
Parameters
----------
path : str or Path
Input file path.
dtype : torch.dtype, optional
Data type for values. If None, inferred from file.
device : str or torch.device
Device to load tensors to.
Returns
-------
SparseTensor
The loaded sparse tensor.
Example
-------
>>> A = load_mtx("matrix.mtx")
>>> A = load_mtx("matrix.mtx", dtype=torch.float32, device="cuda")
"""
from .sparse_tensor import SparseTensor
import numpy as np
path = Path(path)
with open(path, 'r') as f:
# Parse header
header = f.readline().strip()
if not header.startswith("%%MatrixMarket"):
raise ValueError(f"Invalid Matrix Market header: {header}")
parts = header.split()
if len(parts) < 4:
raise ValueError(f"Invalid Matrix Market header: {header}")
obj_type = parts[1].lower() # matrix
format_type = parts[2].lower() # coordinate
field_type = parts[3].lower() if len(parts) > 3 else "real"
symmetry = parts[4].lower() if len(parts) > 4 else "general"
if obj_type != "matrix":
raise ValueError(f"Only 'matrix' object type supported, got: {obj_type}")
if format_type != "coordinate":
raise ValueError(f"Only 'coordinate' format supported, got: {format_type}")
# Skip comments
line = f.readline()
while line.startswith('%'):
line = f.readline()
# Parse dimensions
dims = line.strip().split()
M, N, nnz = int(dims[0]), int(dims[1]), int(dims[2])
# Parse entries
rows = []
cols = []
vals = []
for line in f:
parts = line.strip().split()
if not parts:
continue
r = int(parts[0]) - 1 # Convert to 0-indexed
c = int(parts[1]) - 1
if field_type == "pattern":
v = 1.0
elif field_type == "complex":
v = complex(float(parts[2]), float(parts[3]))
elif field_type == "integer":
v = int(parts[2])
else:
v = float(parts[2])
rows.append(r)
cols.append(c)
vals.append(v)
# Handle symmetry
if symmetry == "symmetric" and r != c:
rows.append(c)
cols.append(r)
vals.append(v)
elif symmetry == "skew-symmetric" and r != c:
rows.append(c)
cols.append(r)
vals.append(-v)
elif symmetry == "hermitian" and r != c:
rows.append(c)
cols.append(r)
vals.append(v.conjugate() if isinstance(v, complex) else v)
# Convert to tensors
row_tensor = torch.tensor(rows, dtype=torch.long, device=device)
col_tensor = torch.tensor(cols, dtype=torch.long, device=device)
# Determine dtype
if dtype is None:
if field_type == "complex":
dtype = torch.complex128
elif field_type == "integer":
dtype = torch.int64
else:
dtype = torch.float64
val_tensor = torch.tensor(vals, dtype=dtype, device=device)
return SparseTensor(val_tensor, row_tensor, col_tensor, (M, N))
[docs]
def load_mtx_info(path: Union[str, Path]) -> Dict:
"""
Read Matrix Market file header without loading data.
Parameters
----------
path : str or Path
Input file path.
Returns
-------
dict
Dictionary with keys: 'shape', 'nnz', 'field', 'symmetry'.
Example
-------
>>> info = load_mtx_info("matrix.mtx")
>>> print(f"Shape: {info['shape']}, NNZ: {info['nnz']}")
"""
path = Path(path)
with open(path, 'r') as f:
header = f.readline().strip()
parts = header.split()
field_type = parts[3].lower() if len(parts) > 3 else "real"
symmetry = parts[4].lower() if len(parts) > 4 else "general"
# Skip comments
line = f.readline()
while line.startswith('%'):
line = f.readline()
dims = line.strip().split()
M, N, nnz = int(dims[0]), int(dims[1]), int(dims[2])
return {
'shape': (M, N),
'nnz': nnz,
'field': field_type,
'symmetry': symmetry,
}
# =============================================================================
# SparseTensor I/O
# =============================================================================
[docs]
def save_sparse(
tensor: "SparseTensor",
path: Union[str, Path],
metadata: Optional[Dict[str, str]] = None
) -> None:
"""
Save a SparseTensor to safetensors format.
Parameters
----------
tensor : SparseTensor
The sparse tensor to save.
path : str or Path
Output file path (should end with .safetensors).
metadata : dict, optional
Additional metadata to store in the file.
Example
-------
>>> A = SparseTensor(val, row, col, (100, 100))
>>> save_sparse(A, "matrix.safetensors")
"""
_ensure_safetensors()
from .sparse_tensor import SparseTensor
if not isinstance(tensor, SparseTensor):
raise TypeError(f"Expected SparseTensor, got {type(tensor)}")
# Prepare tensors dict
tensors = {
"values": tensor.values.contiguous().cpu(),
"row_indices": tensor.row_indices.contiguous().cpu(),
"col_indices": tensor.col_indices.contiguous().cpu(),
"shape": torch.tensor(tensor.sparse_shape, dtype=torch.int64),
}
# Prepare metadata
meta = {
"sparse_dim_0": str(tensor.sparse_dim[0]),
"sparse_dim_1": str(tensor.sparse_dim[1]),
"dtype": str(tensor.dtype),
"format": "sparse_tensor",
"version": "1.0",
}
if metadata:
meta.update(metadata)
save_file(tensors, str(path), metadata=meta)
[docs]
def load_sparse(
path: Union[str, Path],
device: Union[str, torch.device] = "cpu"
) -> "SparseTensor":
"""
Load a SparseTensor from safetensors format.
Parameters
----------
path : str or Path
Input file path.
device : str or torch.device
Device to load tensors to.
Returns
-------
SparseTensor
The loaded sparse tensor.
Example
-------
>>> A = load_sparse("matrix.safetensors", device="cuda")
"""
_ensure_safetensors()
from .sparse_tensor import SparseTensor
tensors = load_file(str(path), device=str(device))
values = tensors["values"]
row_indices = tensors["row_indices"]
col_indices = tensors["col_indices"]
shape = tuple(tensors["shape"].tolist())
return SparseTensor(values, row_indices, col_indices, shape)
# =============================================================================
# Distributed I/O - Save partitioned for multi-rank loading
# =============================================================================
# DSparseTensor persistence -- one safetensors file per rank +
# metadata.json sidecar. Layout v2.0:
# directory/
# metadata.json
# partition_<rank>.safetensors # COO + Partition fields + send/recv idx
_DSPARSE_FORMAT_VERSION = "2.0"
def _placement_axis(placement) -> int:
"""Back-compat axis ID for a VertexShard / VertexShardReplicated /
legacy SparseShard placement, written into metadata.json."""
from .distributed import VertexShard, VertexShardReplicated
if isinstance(placement, VertexShard):
return 0
if isinstance(placement, VertexShardReplicated):
return 1
return int(getattr(placement, "axis", 0))
def _partition_to_safetensors_dict(local_tensor, partition):
out = {
"values": local_tensor.values.contiguous().cpu(),
"row_indices": local_tensor.row_indices.contiguous().cpu(),
"col_indices": local_tensor.col_indices.contiguous().cpu(),
"owned_nodes": partition.owned_nodes.contiguous().cpu(),
"halo_nodes": partition.halo_nodes.contiguous().cpu(),
"local_nodes": partition.local_nodes.contiguous().cpu(),
"global_to_local": partition.global_to_local.contiguous().cpu(),
"local_to_global": partition.local_to_global.contiguous().cpu(),
}
for nid in partition.neighbor_partitions:
if nid in partition.send_indices:
out[f"send_to_{nid}"] = partition.send_indices[nid].contiguous().cpu()
if nid in partition.recv_indices:
out[f"recv_from_{nid}"] = partition.recv_indices[nid].contiguous().cpu()
return out
def _safetensors_dict_to_partition(tensors, neighbor_partitions, partition_id, device):
from .partition import Partition
send_indices, recv_indices = {}, {}
for nid in neighbor_partitions:
if f"send_to_{nid}" in tensors:
send_indices[nid] = tensors[f"send_to_{nid}"].to(device)
if f"recv_from_{nid}" in tensors:
recv_indices[nid] = tensors[f"recv_from_{nid}"].to(device)
return Partition(
partition_id=partition_id,
local_nodes=tensors["local_nodes"].to(device),
owned_nodes=tensors["owned_nodes"].to(device),
halo_nodes=tensors["halo_nodes"].to(device),
neighbor_partitions=list(neighbor_partitions),
send_indices=send_indices,
recv_indices=recv_indices,
global_to_local=tensors["global_to_local"].to(device),
local_to_global=tensors["local_to_global"].to(device),
)
[docs]
def save_dsparse(tensor: "DSparseTensor", directory: Union[str, Path],
rank: Optional[int] = None, verbose: bool = False) -> None:
"""Save this rank's shard. Every rank calls; rank 0 also writes metadata.json."""
_ensure_safetensors()
from .distributed import DSparseTensor
if not isinstance(tensor, DSparseTensor):
raise TypeError(f"Expected DSparseTensor, got {type(tensor).__name__}")
if rank is None:
try:
import torch.distributed as dist
rank = dist.get_rank() if dist.is_initialized() else 0
except (RuntimeError, ImportError):
rank = 0
directory = Path(directory)
directory.mkdir(parents=True, exist_ok=True)
placement = tensor._spec.placement
partition = placement.partition
if partition is None:
raise ValueError("save_dsparse needs a Partition; got Replicated")
save_file(
_partition_to_safetensors_dict(tensor._local_tensor, partition),
str(directory / f"partition_{rank}.safetensors"),
)
if rank == 0:
metadata = {
"format": "dsparse_tensor",
"version": _DSPARSE_FORMAT_VERSION,
"shape": list(tensor.shape),
"dtype": str(tensor.dtype),
"num_partitions": int(tensor.num_partitions),
"placement_axis": _placement_axis(placement),
"partitions": [{
"partition_id": int(partition.partition_id),
"num_owned": int(partition.owned_nodes.numel()),
"num_halo": int(partition.halo_nodes.numel()),
"num_local": int(partition.local_nodes.numel()),
"nnz": int(tensor.nnz),
"neighbors": list(partition.neighbor_partitions),
}],
}
with open(directory / "metadata.json", "w") as f:
json.dump(metadata, f, indent=2)
if verbose:
print(f"[save_dsparse] wrote partition_{rank}.safetensors + metadata.json to {directory}")
elif verbose:
print(f"[save_dsparse] rank {rank}: wrote partition_{rank}.safetensors")
def _resolve_target_world_size(explicit: Optional[int]) -> int:
"""Decide the loader's target world_size.
Precedence: explicit arg β live process group size β 1. Single-
process callers get ``1`` so they can read a sharded archive
without ``torchrun``.
"""
if explicit is not None:
if explicit < 1:
raise ValueError(f"target_world_size must be >= 1, got {explicit}")
return int(explicit)
try:
import torch.distributed as dist
if dist.is_initialized():
return int(dist.get_world_size())
except (RuntimeError, ImportError):
pass
return 1
def _load_all_shards_collapsed(directory: Path,
global_shape: Tuple[int, int],
placement_axis: int,
num_stored: int,
device: Union[str, torch.device]) -> "DSparseTensor":
"""Single-process inverse of ``save_sparse_sharded`` / collective
``save_dsparse``: read every shard, stitch the owned-row slices
back into one global :class:`SparseTensor`, wrap as a
``mesh=None`` (world_size=1) :class:`DSparseTensor` so callers
keep the symmetric API.
Only ``owned`` rows from each shard contribute -- halo rows are
placeholders (zero values) used for halo-exchange addressing on
the saving rank, not real entries.
"""
from .distributed import DSparseTensor
from .partition import Partition
from .sparse_tensor import SparseTensor
device_ = torch.device(device)
all_vals, all_rows, all_cols = [], [], []
for pid in range(num_stored):
local_st, partition = load_sparse_shard(directory, pid, device=device_)
l2g = partition.local_to_global.to(device=device_, dtype=torch.int64)
num_owned = int(partition.owned_nodes.numel())
# Halo rows ([num_owned, num_local)) carry zero placeholder
# entries; drop them before stitching.
owned_mask = local_st.row_indices < num_owned
r_local = local_st.row_indices[owned_mask]
c_local = local_st.col_indices[owned_mask]
all_vals.append(local_st.values[owned_mask])
all_rows.append(l2g[r_local])
all_cols.append(l2g[c_local])
vals = torch.cat(all_vals) if all_vals else torch.empty(0, device=device_)
rows = (torch.cat(all_rows) if all_rows
else torch.empty(0, dtype=torch.int64, device=device_))
cols = (torch.cat(all_cols) if all_cols
else torch.empty(0, dtype=torch.int64, device=device_))
full = SparseTensor(vals, rows, cols, shape=global_shape)
# Trivial single-rank partition: every node owned, no halo,
# local==global. Mirrors ``DSparseTensor`` semantics so the
# returned object behaves as a ``world_size=1`` shard whose
# local tensor IS the global tensor.
n = int(global_shape[0])
arange = torch.arange(n, dtype=torch.int64, device=device_)
trivial = Partition(
partition_id=0,
local_nodes=arange.clone(),
owned_nodes=arange.clone(),
halo_nodes=torch.empty(0, dtype=torch.int64, device=device_),
neighbor_partitions=[],
send_indices={},
recv_indices={},
global_to_local=arange.clone(),
local_to_global=arange.clone(),
)
return DSparseTensor.from_sparse_local(full, mesh=None, partition=trivial,
axis=placement_axis,
global_shape=global_shape)
[docs]
def load_dsparse(directory: Union[str, Path], mesh=None,
rank: Optional[int] = None,
target_world_size: Optional[int] = None,
device: Union[str, torch.device] = "cpu") -> "DSparseTensor":
"""Load a sharded archive into a :class:`DSparseTensor`.
Three modes, picked from ``target_world_size`` (or, if omitted,
a live process group, else ``1``):
* ``stored_N == target_world_size`` -- fast path; this rank reads
``partition_{rank}.safetensors`` directly. Original
:func:`save_dsparse` semantics.
* ``target_world_size == 1`` -- single-process gather: read every
shard from disk, stitch the owned rows into one global
:class:`SparseTensor`, return as a ``mesh=None`` trivial
DSparseTensor. Useful for inspection / single-node debug after
sharded training.
* otherwise -- explicit :class:`NotImplementedError`. True
re-partitioning on load (``stored_N != target != 1``) requires
a cross-rank shuffle; deferred to a future ``redistribute()``
pass. Workaround: load with ``target_world_size=1`` and call
``.full_tensor()`` (or save again at the new world size).
"""
_ensure_safetensors()
from .distributed import DSparseTensor
from .sparse_tensor import SparseTensor
directory = Path(directory)
metadata = load_metadata(directory)
if metadata.get("format") != "dsparse_tensor":
raise ValueError(f"not a DSparseTensor directory: {directory}")
global_shape = tuple(metadata["shape"])
placement_axis = int(metadata.get("placement_axis", 0))
num_stored = int(metadata["num_partitions"])
target_N = _resolve_target_world_size(target_world_size)
if num_stored != target_N:
if target_N == 1:
return _load_all_shards_collapsed(
directory, global_shape, placement_axis, num_stored, device)
raise NotImplementedError(
f"stored {num_stored} shards != target world_size {target_N}; "
"in-place repartitioning is not implemented yet. "
"Workaround: load with target_world_size=1, call "
".full_tensor(), then save_sparse_sharded() at the new size."
)
if rank is None:
try:
import torch.distributed as dist
rank = dist.get_rank() if dist.is_initialized() else 0
except (RuntimeError, ImportError):
rank = 0
shard_path = directory / f"partition_{rank}.safetensors"
if not shard_path.exists():
raise FileNotFoundError(f"missing shard for rank {rank}: {shard_path}")
tensors = load_file(str(shard_path), device=str(device))
# Neighbor list reconstructed from safetensors keys -- robust even
# if metadata.partitions was only written by rank 0.
neighbors = sorted({int(k.removeprefix("send_to_")) for k in tensors if k.startswith("send_to_")}
| {int(k.removeprefix("recv_from_")) for k in tensors if k.startswith("recv_from_")})
partition = _safetensors_dict_to_partition(tensors, neighbors, rank, device)
n_local = int(partition.local_nodes.numel())
local_st = SparseTensor(tensors["values"], tensors["row_indices"],
tensors["col_indices"], shape=(n_local, n_local))
if mesh is None:
import torch.distributed as dist
if dist.is_initialized():
try:
from torch.distributed.device_mesh import init_device_mesh
except ImportError:
from torch.distributed._tensor.device_mesh import init_device_mesh
mesh = init_device_mesh(torch.device(device).type, (dist.get_world_size(),))
return DSparseTensor.from_sparse_local(local_st, mesh, partition,
axis=placement_axis,
global_shape=global_shape)
def save_sparse_sharded(tensor: "SparseTensor", directory: Union[str, Path],
num_partitions: int, partition_method: str = "simple",
coords: Optional[torch.Tensor] = None,
verbose: bool = False) -> None:
"""Single-process partition + write all shards. Output is identical
to a collective ``save_dsparse`` and can be consumed under torchrun."""
_ensure_safetensors()
from .partition import resolve_partition_ids, build_partition
if tensor.is_batched:
raise ValueError("save_sparse_sharded does not support batched SparseTensor")
directory = Path(directory)
directory.mkdir(parents=True, exist_ok=True)
partition_ids = resolve_partition_ids(
tensor.row_indices, tensor.col_indices, int(tensor.shape[0]),
num_partitions, method=partition_method, coords=coords,
)
per_partition_meta = []
for pid in range(num_partitions):
partition = build_partition(
tensor.row_indices, tensor.col_indices, int(tensor.shape[0]),
partition_ids.cpu(), pid,
)
local_st = tensor.extract_partition(partition)
save_file(_partition_to_safetensors_dict(local_st, partition),
str(directory / f"partition_{pid}.safetensors"))
per_partition_meta.append({
"partition_id": int(partition.partition_id),
"num_owned": int(partition.owned_nodes.numel()),
"num_halo": int(partition.halo_nodes.numel()),
"num_local": int(partition.local_nodes.numel()),
"nnz": int(local_st.nnz),
"neighbors": list(partition.neighbor_partitions),
})
metadata = {
"format": "dsparse_tensor",
"version": _DSPARSE_FORMAT_VERSION,
"shape": list(tensor.sparse_shape),
"dtype": str(tensor.dtype),
"num_partitions": int(num_partitions),
"placement_axis": 0,
"partition_method": partition_method,
"partitions": per_partition_meta,
}
with open(directory / "metadata.json", "w") as f:
json.dump(metadata, f, indent=2)
if verbose:
print(f"[save_sparse_sharded] wrote {num_partitions} partitions + metadata.json to {directory}")
def load_sparse_shard(directory: Union[str, Path], rank: int,
device: Union[str, torch.device] = "cpu") -> Tuple["SparseTensor", "object"]:
"""Load one shard for inspection. Returns ``(local_tensor, partition)``."""
_ensure_safetensors()
from .sparse_tensor import SparseTensor
directory = Path(directory)
metadata = load_metadata(directory)
if metadata.get("format") != "dsparse_tensor":
raise ValueError(f"not a DSparseTensor directory: {directory}")
shard_path = directory / f"partition_{rank}.safetensors"
if not shard_path.exists():
raise FileNotFoundError(f"missing shard for rank {rank}: {shard_path}")
tensors = load_file(str(shard_path), device=str(device))
neighbors = sorted({int(k.removeprefix("send_to_")) for k in tensors if k.startswith("send_to_")}
| {int(k.removeprefix("recv_from_")) for k in tensors if k.startswith("recv_from_")})
partition = _safetensors_dict_to_partition(tensors, neighbors, rank, device)
n_local = int(partition.local_nodes.numel())
local_st = SparseTensor(tensors["values"], tensors["row_indices"],
tensors["col_indices"], shape=(n_local, n_local))
return local_st, partition