public func smoothL1Loss(
    predictions: MLXArray, targets: MLXArray, beta: Float = 1, reduction: LossReduction = .mean
) -> MLXArray {
    precondition(predictions.shape == targets.shape)

    let diff = abs(predictions - targets)
    let loss = which(diff .< beta, 0.5 * square(diff) / beta, abs(diff) - 0.5 * beta)

    return reduction.reduce(loss: loss)
}

    let diff = predictions - targets

    let diff = abs(predictions - targets)