ls_mlkit.util.shape_class module

class ls_mlkit.util.shape_class.Shape(config: ShapeConfig)[source]

Bases: object

complete_micro_shape(x: Tensor) → Tensor[source]

Complete the micro shape of \(x\), assuming the macro shape is already known

Parameters:

x (Tensor) – \(x\)

Returns:

\(x\) with the micro shape completed

Return type:

Tensor

get_macro_shape(x: Tensor) → tuple[int, ...][source]

Get the macro shape of \(x\)

Parameters:

x (Tensor) – \(x\)

Returns:

the shape of the macro part of \(x\)

Return type:

tuple[int, ...]

class ls_mlkit.util.shape_class.ShapeConfig(ndim_micro_shape: int)[source]

Bases: object