DeepSeek V4.1

def make_identity_pre_mix(x: torch.Tensor, hc_mult: int) -> torch.Tensor:
    """initial one-hot mix"""
    pre_mix = x.new_zeros(x.size(0), x.size(1), hc_mult, dtype=torch.float32)
    pre_mix[:, :, 0] = 1.0
    return pre_mix

    @torch.inference_mode()
    def forward(
        self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None
    ):
        """input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it
        has already filled, so the padding it uses internally never reaches here. token_types /
        images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie
        inside the first (start_pos 0) chunk."""
        image_mask = None if token_types is None else token_types >= 0  # TEXT is -1
        # image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask
        engram_mask = None if image_mask is None else ~image_mask
        engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None
        h = self.embed(input_ids)
        if images is not None:
            assert start_pos == 0, "image spans must be prefilled in a single chunk"
            self.merge_image_embeddings(images, h)
        # Expand to hc_mult copies for Hyper-Connections
        h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
        main_hiddens = []
        pre_mix = make_identity_pre_mix(h, self.hc_mult)
        for i, layer in enumerate(self.layers):
            if layer.engram is not None:
                h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask)
            # the MTP head reads the attention input of its target layers, not their output
            if i in self.target_layer_ids:
                main_hiddens.append(h.mean(dim=2))
            h, pre_mix = layer(h, start_pos, pre_mix, image_mask)
        h = layer.hc_pre(h, pre_mix)
        logits = self.head(self.norm(h))
        output_ids = sample(logits, self.temperature)
        main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None
        return output_ids, logits, main_hidden