Artifact builder¶

class torch_lattice.artifact.builder.ModuleLowering(fn, arities)[source]¶

Bases: object

Parameters:

fn (Callable[[...], ArtifactValue])
arities (frozenset[int])

fn: Callable[[...], ArtifactValue]¶

arities: frozenset[int]¶

lower(builder, name, module, inputs)[source]¶

Return type:

ArtifactValue

Parameters:

builder (TorchLatticeArtifactBuilder)
name (str)
module (Module)
inputs (tuple[ArtifactValue, ...])

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

Register a Torch module lowering method.

Return type:

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

Parameters:

types (type[Module])
arity (int | tuple[int, ...])

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:

name (str)
module (Module)

lower_module(name, module, *inputs)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Module)
inputs (ArtifactValue)

conv3d(name, module, input, coordinates=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Conv3d)
input (ArtifactValue)
coordinates (ArtifactValue | None)

subm_conv3d(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (SubmConv3d)
input (ArtifactValue)

normalized_subm_conv3d(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (NormalizedSubmConv3d)
input (ArtifactValue)

conv_transpose3d(name, module, input, coordinates=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (ConvTranspose3d)
input (ArtifactValue)
coordinates (ArtifactValue | None)

normalized_conv_transpose3d(name, module, input, coordinates=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (NormalizedConvTranspose3d)
input (ArtifactValue)
coordinates (ArtifactValue | None)

generative_conv_transpose3d(name, module, input, coordinates=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (GenerativeConvTranspose3d)
input (ArtifactValue)
coordinates (ArtifactValue | None)

normalized_generative_conv_transpose3d(name, module, input, coordinates=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (NormalizedGenerativeConvTranspose3d)
input (ArtifactValue)
coordinates (ArtifactValue | None)

pool3d(name, module, input=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Pool3d)
input (ArtifactValue | None)

pool_transpose3d(name, module, input, target=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (PoolTranspose3d)
input (ArtifactValue)
target (ArtifactValue | None)

trilinear_upsample3d(name, module, input, target=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (TrilinearUpsample3d)
input (ArtifactValue)
target (ArtifactValue | None)

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

Return type:

ArtifactValue

Parameters:

name (str)
points (ArtifactValue)
features (ArtifactValue)
batch_indices (ArtifactValue)
active_rows (ArtifactValue)
reduction (Literal['sum', 'mean'])

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

Return type:

ArtifactValue

Parameters:

name (str)
points (ArtifactValue)
voxels (ArtifactValue)
batch_indices (ArtifactValue)
point_active_rows (ArtifactValue)
interpolation (Literal['nearest', 'linear'])

batch_norm(name, module, input=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (BatchNorm)
input (ArtifactValue | None)

layer_norm(name, module, input=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (LayerNorm)
input (ArtifactValue | None)

rms_norm(name, module, input=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Module)
input (ArtifactValue | None)

relu(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Module)
input (ArtifactValue)

leaky_relu(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (LeakyReLU | LeakyReLU)
input (ArtifactValue)

silu(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Module)
input (ArtifactValue)

gelu(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (GELU | GELU)
input (ArtifactValue)

sigmoid(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Module)
input (ArtifactValue)

tanh(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Module)
input (ArtifactValue)

softplus(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Softplus | Softplus)
input (ArtifactValue)

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

Return type:

ArtifactValue

Parameters:

name (str)
kind (str)
input (ArtifactValue | None)
approximate (str)
alpha (float)
beta (float)
threshold (float)

global_avg_pool(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (GlobalAvgPool)
input (ArtifactValue)

global_max_pool(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (GlobalMaxPool)
input (ArtifactValue)

global_sum_pool(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (GlobalSumPool)
input (ArtifactValue)

global_pool(name, mode, input=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
mode (Literal['sum', 'avg', 'max'])
input (ArtifactValue | None)

linear(name, module, input=None)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Linear)
input (ArtifactValue | None)

dropout(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Dropout)
input (ArtifactValue)

flatten(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Flatten)
input (ArtifactValue)

identity(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (Identity)
input (ArtifactValue)

instance_norm(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (InstanceNorm)
input (ArtifactValue)

group_norm(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (GroupNorm)
input (ArtifactValue)

sparse_crop(name, module, input)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
module (SparseCrop)
input (ArtifactValue)

sparse_binary(name, lhs, rhs, op, *, join='outer', lhs_fill=0.0, rhs_fill=0.0)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
lhs (ArtifactValue)
rhs (ArtifactValue)
op (str)
join (str)
lhs_fill (float)
rhs_fill (float)

sparse_add(name, lhs, rhs, *, join='outer', lhs_fill=0.0, rhs_fill=0.0)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
lhs (ArtifactValue)
rhs (ArtifactValue)
join (str)
lhs_fill (float)
rhs_fill (float)

sparse_cat(name, lhs, rhs, *, join='inner')[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
lhs (ArtifactValue)
rhs (ArtifactValue)
join (str)

sparse_reindex(name, input, target, *, fill=0.0)[source]¶

Return type:

ArtifactValue

Parameters:

name (str)
input (ArtifactValue)
target (ArtifactValue)
fill (float)

output(*values, names=None)[source]¶

Return type:

Parameters:

values (ArtifactValue)
names (tuple[str, ...] | None)

to_mlir()[source]¶

Return type:: str

save(artifact_dir, *, clean=True, validate=True)[source]¶

Parameters:

artifact_dir (str | Path)
clean (bool)
validate (bool)