How to get the logits in mlx-lm?

I'm interested in adding mlx-lm as a provider in Backtrack Sampler.

I generate one token at a time, apply my own temperature, and handle the selection of the next token in my own code. Therefore, I need to retrieve the logits from mlx-lm rather than pass it a custom sampler.

I have some changes in a branch, fwiw, but they currently don't work because I handle log probabilities (logprobs) instead of logits, as I'm using the utils.stream_generate function. I don't necessarily need to use stream_generate; I used it because I saw it here.

Is there a way to get the logits in mlx-lm?

Start by reading the mlx-lm utils.stream_generate entry point and the linked Backtrack Sampler feature branch. Trace how log probabilities are produced and determine whether the integration can retrieve logits for one-token-at-a-time selection; done means the provider can apply its own temperature and choose the next token from logits.