Mila::Dnn::Serialization::SafeTensorsWriter Class Reference

Writes a safetensors file: u64 header length, JSON header, packed data. More...

Public Member Functions
	SafeTensorsWriter (const SafeTensorsWriter &)=delete
	SafeTensorsWriter (const std::filesystem::path &filepath)
	Open a path for writing.
void	beginData ()
	Emit the header and open the data region.
void	close ()
	Flush and close.
void	declareTensor (const std::string &name, TensorDataType dtype, const shape_t &shape)
	Declare a tensor and reserve its byte range.
size_t	getTensorCount () const noexcept
	Number of declared tensors.
SafeTensorsWriter &	operator= (const SafeTensorsWriter &)=delete
void	setMetadata (const std::string &key, const std::string &value)
	Set a metadata entry.
void	writeTensorData (const std::string &name, const void *data, size_t nbytes)
	Append one tensor body.

Detailed Description

Writes a safetensors file: u64 header length, JSON header, packed data.

Two-phase by necessity. The header records every tensor's byte range, so all shapes must be known before any data is written; declare everything, then stream the bodies in declaration order. That ordering is not a convenience – the container requires the data region to be tiled contiguously, so a body written out of order corrupts the file. Streaming rather than buffering is what lets a 22 GB model be written while only one staged tensor is resident.

Example

SafeTensorsWriter writer( path );
writer.declareTensor( "weight", TensorDataType::BF16, shape_t{ 4096, 4096 } );
writer.setMetadata( kMilaConfigMetadataKey, config_json.dump() );
writer.beginData();
writer.writeTensorData( "weight", host_bytes, byte_count );
writer.close();

Constructor & Destructor Documentation

SafeTensorsWriter()

Mila::Dnn::Serialization::SafeTensorsWriter::SafeTensorsWriter ( const std::filesystem::path & filepath )

inlineexplicit

Open a path for writing.

Exceptions

std::runtime_error

if the file cannot be created.

Member Function Documentation

beginData()

void Mila::Dnn::Serialization::SafeTensorsWriter::beginData ( )

inline

Emit the header and open the data region.

Ends phase one.

close()

void Mila::Dnn::Serialization::SafeTensorsWriter::close ( )

inline

Flush and close.

Exceptions

std::runtime_error

if any declared tensor was never written – the header would promise bytes the file does not contain.

declareTensor()

void Mila::Dnn::Serialization::SafeTensorsWriter::declareTensor ( const std::string & name, TensorDataType dtype, const shape_t & shape )

inline

Declare a tensor and reserve its byte range.

Phase one.

Exceptions

std::runtime_error

if data writing has begun, or the name repeats.

setMetadata()

void Mila::Dnn::Serialization::SafeTensorsWriter::setMetadata ( const std::string & key, const std::string & value )

inline

Set a metadata entry.

Phase one.

The container permits string values only; encode structured data as JSON text.

writeTensorData()

void Mila::Dnn::Serialization::SafeTensorsWriter::writeTensorData ( const std::string & name, const void * data, size_t nbytes )

inline

Append one tensor body.

Phase two, in declaration order.

Exceptions

std::runtime_error

on order or size mismatch, either of which would silently desynchronize the file from its own index.

---

The documentation for this class was generated from the following file:

• Mila/Src/Dnn/Serialization/SafeTensors.ixx