fix(types): improve type hints for mlx_lm.utils.load method

[rahuliyer95]

Problem

The mlx_lm.utils.load returns a Tuple[nn.Module, TokenizerWrapper] when return_config=False and Tuple[nn.Module, TokenizerWrapper, Dict[str, Any]] when return_config=True. Since the return type is controlled by a function argument, the types can improved to determine the correct return type based on the value of return_config.

Solution

We can use the @overload decroator to help provide the proper type hints to the callers based on the value of return_config.

[awni]

Thanks but it seems unnecessary. Also would prefer not to reorder the positions of the arguments just in case...

[rahuliyer95]

Thanks but it seems unnecessary. Also would prefer not to reorder the positions of the arguments just in case...

I understand the concern about reordering the positions of the arguments, but using mlx-lm in a project with mypy complains about wrong return type which can be easily avoided with this change.

[awni]

Hmm.. ok fair enough. Will you update the PR to use the old ordering and then we can merge it?

[angeloskath]

Just to be clear, mypy complains because the return type is a Union and you don't somehow verify that it is one of the two cases in the union when you pass it to a different function right?

It feels like the issue is not in mlx-lm and even though the @overload would help mypy deduce the type better I think the return type is fine and don't really like having the function definition replicated 3 times for little benefit...

Wdyt @rahuliyer95 ?

[rahuliyer95]

Just to be clear, mypy complains because the return type is a Union and you don't somehow verify that it is one of the two cases in the union when you pass it to a different function right?

Yes that's correct.

It feels like the issue is not in mlx-lm and even though the @overload would help mypy deduce the type better I think the return type is fine and don't really like having the function definition replicated 3 times for little benefit...

Wdyt @rahuliyer95 ?

I believe the benefit is high because without this every consumer of mlx-lm would have to have a check in their code which can be avoided with this overload. Here is some prior art I found for other libraries like Tranformers, PyTorch which do the same.