Artifact builder

class torch_lattice.artifact.builder.ModuleLowering(types, fn, arity)[source]

Bases: object

Parameters:
types: tuple[type[Module], ...]
fn: Callable[[...], ArtifactValue]
arity: int
matches(module)[source]
Return type:

bool

Parameters:

module (Module)

lower(builder, name, module, inputs)[source]
Return type:

ArtifactValue

Parameters:
torch_lattice.artifact.builder.module_lowering(*types, arity=1)[source]

Register a Torch module lowering method.

Return type:

Callable[[Callable[..., ArtifactValue]], Callable[..., ArtifactValue]]

Parameters:
class torch_lattice.artifact.builder.ArtifactValue(value, kind, channels)[source]

Bases: object

A value in the Torch-to-lattice artifact graph.

Parameters:

  • value (object)

  • kind (Literal['sparse_tensor', 'dense_tensor'])

  • channels (int | None)

value: object
kind: Literal['sparse_tensor', 'dense_tensor']
channels: int | None
class torch_lattice.artifact.builder.TorchLatticeArtifactBuilder(*, input_name='input', output_name='output', input_dtype='f32', batch_size=None, quantize_bits=None, quantize_group_size=32, quantize_scale_dtype='f16', create_default_input=True, input_stride=(1, 1, 1))[source]

Bases: object

Explicit builder for Torch-to-lattice MLIR artifacts.

Parameters:

  • input_name (str)

  • output_name (str)

  • input_dtype (str)

  • batch_size (int | None)

  • quantize_bits (int | None)

  • quantize_group_size (int)

  • quantize_scale_dtype (str)

  • create_default_input (bool)

property current: ArtifactValue
property weights: dict[str, Tensor]
sparse_input()[source]
Return type:

ArtifactValue

sparse_argument(name, *, dtype=None, channels=None, stride=(1, 1, 1))[source]
Return type:

ArtifactValue

Parameters:

  • name (str)

  • dtype (str | None)

  • channels (int | None)

dense_argument(name, type, *, channels=None)[source]
Return type:

ArtifactValue

Parameters:

  • name (str)

  • type (str | TensorType)

  • channels (int | None)

module(name, module)[source]
Return type:

ArtifactValue

Parameters:
lower_module(name, module, *inputs)[source]
Return type:

ArtifactValue

Parameters:
conv3d(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
subm_conv3d(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
conv_transpose3d(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
generative_conv_transpose3d(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
target_conv3d(name, module, input, target)[source]
Return type:

ArtifactValue

Parameters:
pool3d(name, module, input=None)[source]
Return type:

ArtifactValue

Parameters:
voxelize(name, *, points, features, batch_indices, active_rows, voxel_size, origin=0.0, reduction='mean', stride=1)[source]
Return type:

ArtifactValue

Parameters:
devoxelize(name, *, points, voxels, batch_indices, point_active_rows, voxel_size, origin=0.0, interpolation='nearest')[source]
Return type:

ArtifactValue

Parameters:
batch_norm(name, module, input=None)[source]
Return type:

ArtifactValue

Parameters:
layer_norm(name, module, input=None)[source]
Return type:

ArtifactValue

Parameters:
rms_norm(name, module, input=None)[source]
Return type:

ArtifactValue

Parameters:
relu(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
leaky_relu(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
silu(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
gelu(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
sigmoid(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
tanh(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
softplus(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
activation(name, kind, *, input=None, approximate='none', alpha=0.01, beta=1.0, threshold=20.0)[source]
Return type:

ArtifactValue

Parameters:
global_avg_pool(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
global_max_pool(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
global_pool(name, mode, input=None)[source]
Return type:

ArtifactValue

Parameters:
linear(name, module, input=None)[source]
Return type:

ArtifactValue

Parameters:
dropout(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
flatten(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
identity(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
instance_norm(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
group_norm(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
sparse_crop(name, module, input)[source]
Return type:

ArtifactValue

Parameters:
sparse_binary(name, lhs, rhs, op, *, join='outer', lhs_fill=0.0, rhs_fill=0.0)[source]
Return type:

ArtifactValue

Parameters:
sparse_add(name, lhs, rhs, *, join='outer', lhs_fill=0.0, rhs_fill=0.0)[source]
Return type:

ArtifactValue

Parameters:
sparse_cat(name, lhs, rhs, *, join='inner')[source]
Return type:

ArtifactValue

Parameters:
output(value=None, *, name=None)[source]
Return type:

None

Parameters:
to_mlir()[source]
Return type:

str

save(artifact_dir, *, clean=True, validate=True)[source]
Parameters:
class torch_lattice.artifact.builder.PackedWeight(weight, scales, biases)[source]

Bases: object

Parameters:
weight: Tensor
scales: Tensor
biases: Tensor
torch_lattice.artifact.builder.dequantize_artifact_weight(tensor, *, bits, group_size, scale_dtype='f16')[source]

Return the logical weight represented by artifact quantization.

This uses the same packing path as artifact export and then unpacks the affine integer payload back to the source tensor layout. Test fixture expected outputs can use it to compare deployment semantics instead of dense pre-quantization training semantics.

Return type:

Tensor

Parameters: