Skip to main content

clika_runtime.nn functions

fuse_convs​

fuse_convs(*layers: 'Conv', name: 'str | None' = None, activation: 'ActivationLike' = None) -> 'Conv'

fuse_convs(*layers, name=None, *, activation=None) -> Conv

Fuse layers (two or more :class:~clika_runtime.nn.Conv layers, the Conv1d / Conv2d / Conv3d spellings included, over the same input, bound and not yet run) into one convolution over their concatenated weights, and return the group as a :class:~clika_runtime.nn.Conv: its forward maps [batch, *spatial, in_channels] to [batch, *spatial', sum(out_channels)], the parts' outputs concatenated along the channel axis in the order given (split it with :func:~clika_runtime.split_with_sizes).

Parts share one convolution when their geometry agrees: the kernel, the stride, the dilation, the padding and its mode, groups=1, one dtype, and a bias on every part or on none. A group the runtime cannot fuse (another kernel on the same grid, grouped convolutions, mixed biases) is returned in the same shape, serving the parts side by side plus the concatenation: the output never depends on whether the fusion engaged, only the dispatch count does.

The parts keep their state-dict names and their weight and bias views, each reading back its channels of the fused weight and bias, so state_dict() of the model stays the checkpoint's and the group's pack is the one resident copy; the group itself lists no parameters. After the fuse the parts no longer serve a forward of their own: calling a part raises :class:~clika_runtime.ClikaRTError until a load_state_dict over it dissolves the group (the group then serves the parts side by side plus the concatenation).

activation fuses a pointwise epilogue into the group's forward. Parts that each apply the same activation keep it: the group applies it once over the concatenation, and activation may repeat it but never name another. name labels the group in its repr.

Raises: TypeError: when an argument is not a :class:~clika_runtime.nn.Conv. ValueError: when fewer than two layers are given. ClikaRTError: when a part is unbound or already packed (fuse before the first forward), the parts' own activations disagree, or activation is a gated one (it halves the channels, which a convolution does not serve).

fuse_linears​

fuse_linears(*layers: 'Linear | QLinearWoQ', name: 'str | None' = None, activation: 'ActivationLike' = None) -> 'Linear'

fuse_linears(*layers, name=None, *, activation=None) -> Linear

Fuse layers (two or more :class:~clika_runtime.nn.Linear or :class:~clika_runtime.nn.QLinearWoQ layers over the same input width, bound and not yet run) into one matmul over their concatenated weights, and return the group as a :class:~clika_runtime.nn.Linear: its forward maps [*, in_features] to [*, sum(out_features)], the parts' outputs concatenated in the order given (split it with :func:~clika_runtime.split_with_sizes).

The parts keep their state-dict names and their weight views, each reading back its rows of the fused weight (a view held from before the fuse becomes storage-free and answers its metadata; part.weight reconstructs the rows on demand), so state_dict() of the model stays the checkpoint's and the group holds the one resident copy of every fused weight. The group serves at its parts' residency when they agree and packed when they differ, so under the "per_call" residency that copy is the joined weight itself rather than a pack. The group itself lists no parameters (group.state_dict() is empty), so holding it as a module attribute adds nothing to a save. After the fuse the parts no longer serve a forward of their own: calling a part raises :class:~clika_runtime.ClikaRTError until a load_state_dict over it dissolves the group (the group then serves the parts side by side plus the concatenation). A group the runtime cannot fuse (mixed quantization schemes, unequal input widths) is returned in that same shape: the output never depends on whether the fusion engaged, only the dispatch count does. Weight-quantized parts share one matmul only under one quantization scheme; parts under different schemes serve side by side.

activation fuses an epilogue into the group's forward; a gated one ("swiglu", "geglu", "reglu", "situ") emits half the concatenated width, the gate / up convention, and needs bias-free parts. Parts that each apply the same pointwise activation keep it: the group applies it once over the concatenation, and activation may repeat it but never name another. name labels the group in its repr.

Raises: TypeError: when an argument is not a :class:~clika_runtime.nn.Linear or :class:~clika_runtime.nn.QLinearWoQ (a :class:~clika_runtime.nn.QLinear takes a quantized input, which a group does not). ValueError: when fewer than two layers are given. ClikaRTError: when a part is unbound or already packed (fuse before the first forward), the parts carry biases under a gated activation, or the parts' own activations disagree (parts that apply different ones, a gated one on a part, or an activation that names another).