Xcode warning in `applyRepetitionPenalty`

[Anthony DePasquale]

When compiling the LLMEval example app, I get a warning from Xcode about the variable logits. Sorry if this is too basic a question or I'm missing something, but I believe in Swift classes are reference types, but var logits = logits implies that a copy is being created. Isn't this function simply returning a reference to logits, which it mutates?

private func applyRepetitionPenalty(logits: MLXArray, repetitionContext: MLXArray, penalty: Float) -> MLXArray {
    var logits = logits // Xcode warning: Variable 'logits' was never mutated; consider changing to 'let' constant

    if repetitionContext.shape[0] > 0 {
        let indices = repetitionContext
        var selectedLogits = take(logits, indices, axis: -1).squeezed(axis: 0)

        selectedLogits = MLX.where(
            selectedLogits .< 0, selectedLogits * penalty, selectedLogits / penalty)

        logits[0..., indices] = selectedLogits
        return logits
    }

    return logits
}

[David Koski]

Right, it is a legit warning. logits is a reference to an instance of MLXArray, which is a class (reference type). There is no assignment to logits, that is the reference never changes, so it can be written as a let.

This line:

        logits[0..., indices] = selectedLogits

appears to mutate it, but it actually mutates the inner state of logits and isn't really an assignment.

Here is a case where an MLXArray variable has to be a var:

    func callAsFunction(_ x: MLXArray) -> MLXArray {
        var x = x
        
        for resnet in resnets {
            x = resnet(x)
        }
        
        if let downsample {
            x = padded(x, widths: [[0, 0], [0, 1], [0, 1], [0, 0]])
            x = downsample(x)
        }
        
        if let upsample {
            x = upsample(upsampleNearest(x))
        }
        
        return x
    }

In particular the x = are all assignments so the the reference itself is being mutated.

Make sense?

[Anthony DePasquale]

Yes, that makes sense. Is there any reason then to declare let logits = logits in applyRepetitionPenalty? It seems like this could be omitted.

[David Koski]

Agreed, it is not needed. I think in a past version it must have reassigned it and this is what is left over.