Method forward

exir/tests/test_memory_planning.py:1353–1360

(
                self,
                input_pos: torch.Tensor,
                k_val: torch.Tensor,
                v_val: torch.Tensor,
            )

                )

            def forward(
                self,
                input_pos: torch.Tensor,
                k_val: torch.Tensor,
                v_val: torch.Tensor,
            ) -> torch.Tensor:
                k_out, v_out = self.kv_cache.update(input_pos, k_val, v_val)
                return (k_out + v_out).sum(dim=-1)

            def reset(self, k_zeros: torch.Tensor, v_zeros: torch.Tensor) -> None:
                self.kv_cache.k_cache.copy_(k_zeros)