Correct the type annotation of cache in llama.py

[wangkuiyi]

The source code accesses cache.offset and cache.update_and_fetch, which belong to base.KVCache, so it seems that we should update the type annotation of cache from Tuple[mx.array, mx.array] into base.KVCache.

After doing so, mypy does not longer complain about the type annotation.

[awni]

Btw @wangkuiyi you can add the pre-commit hooks it will run all the formatting automatically: https://github.com/ml-explore/mlx-examples/blob/main/CONTRIBUTING.md#pull-requests

[wangkuiyi]

Btw @wangkuiyi you can add the pre-commit hooks it will run all the formatting automatically: https://github.com/ml-explore/mlx-examples/blob/main/CONTRIBUTING.md#pull-requests

Yeah. I just noticed that CircleCI calls pre-commit. Would it be useful if I start to fix mypy errors for now before enabling mypy in pre-commit? @awni

[awni]

Would it be useful if I start to fix mypy errors for now before enabling mypy in pre-commit?

For sure. We'll gladly pull in any type annotation fixes if you send them.