Artifact IO

class torch_lattice.artifact.io.LatticeArtifactSaveResult(artifact_dir, graph_path, weights_path, weight_keys)[source]

Bases: object

Result of exporting a Torch model to a lattice artifact directory.

Parameters:
artifact_dir: Path
graph_path: Path
weights_path: Path
weight_keys: tuple[str, ...]
exception torch_lattice.artifact.io.LatticeModelArtifactError[source]

Bases: ValueError

class torch_lattice.artifact.io.LatticeModelArtifactOptions(input_dtype='f32', batch_size=None, clean=True, validate=True, quantize_bits=None, quantize_group_size=32, quantize_scale_dtype='f16', input_stride=(1, 1, 1))[source]

Bases: object

Options for producing a portable lattice MLIR artifact.

Parameters:

  • input_dtype (str)

  • batch_size (int | None)

  • clean (bool)

  • validate (bool)

  • quantize_bits (int | None)

  • quantize_group_size (int)

  • quantize_scale_dtype (str)

  • input_stride (tuple[int, int, int])

input_dtype: str = 'f32'
batch_size: int | None = None
clean: bool = True
validate: bool = True
quantize_bits: int | None = None
quantize_group_size: int = 32
quantize_scale_dtype: str = 'f16'
input_stride: tuple[int, int, int] = (1, 1, 1)
torch_lattice.artifact.io.save_lattice_model_artifact(model, artifact_dir, *, input_name='input', output_name='output', sample_input=None, options=None, method='fx')[source]

Save model as graph.mlir plus weights.safetensors.

torch.fx is the default front-end. For explicit construction use TorchLatticeArtifactBuilder directly and call save.

Return type:

LatticeArtifactSaveResult

Parameters: