def forward(
        self,
        tokens: Tensor,
        H: HiddenState | None = None,
        pad: Tensor | None = None,
        hooks: TransformerHooks = TransformerHooks(),
    ) -> tuple[Tensor, HiddenState]:
        """
        Forward pass through the transformer!

        During evaluation or first forward pass in sampling:
            X is expected to be a [batch_size x sequence_length]-shaped LongTensor of encoded prompts.
            H is expected to be None.
            pad is a [batch_size x sequence_length]-shaped boolean Tensor. "1"s mean "ignore this
            token". This parameter must be set if not all encoded prompts in X have the same length.
            Note that activations observed by hooks will include padded values.

        During sampling after first forward pass:
            X is expected to be the new part of the sequences (eg most recently sampled tokens).
            H is expected to have KV-caches of all Keys and Values for prior tokens.
            pad is expected to be None (new tokens are not pad tokens).

        Returns a tuple containing the resulting logits tensor and a new hidden state consisting of a KV cache.
        """
        X, H, pad, hooks = self.run_embed(tokens, H, pad, hooks)
        X, H, pad, hooks = self.run_torso(X, H, pad, hooks)
        return self.run_unembed(X, H, hooks)

        Forward pass through the transformer!

        During evaluation or first forward pass in sampling:
            X is expected to be a [batch_size x sequence_length]-shaped LongTensor of encoded prompts.
            H is expected to be None.
            pad is a [batch_size x sequence_length]-shaped boolean Tensor. "1"s mean "ignore this
            token". This parameter must be set if not all encoded prompts in X have the same length.
            Note that activations observed by hooks will include padded values.

        During sampling after first forward pass:
            X is expected to be the new part of the sequences (eg most recently sampled tokens).
            H is expected to have KV-caches of all Keys and Values for prior tokens.
            pad is expected to be None (new tokens are not pad tokens).

        Returns a tuple containing the resulting logits tensor and a new hidden state consisting of a KV cache.
        
neuron_explainer/models/transformer.py