diff --git a/swift/Sources/CoreAIImageSegmenter/Postprocessing/SegmentationPostprocessor.swift b/swift/Sources/CoreAIImageSegmenter/Postprocessing/SegmentationPostprocessor.swift index 7e2ab99a..d86dae73 100644 --- a/swift/Sources/CoreAIImageSegmenter/Postprocessing/SegmentationPostprocessor.swift +++ b/swift/Sources/CoreAIImageSegmenter/Postprocessing/SegmentationPostprocessor.swift @@ -3,6 +3,7 @@ // Use of this source code is governed by a BSD-3-clause license that can // be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause +import CoreAIShared import CoreGraphics import Foundation @@ -227,11 +228,6 @@ public enum SegmentationPostprocessor { probabilities: probabilities, width: outputWidth, height: outputHeight) } - // MARK: - Helpers - public static func sigmoid(_ x: Float) -> Float { - 1.0 / (1.0 + exp(-x)) - } - /// Bilinear-upsample a Float grid and threshold at the destination resolution. /// /// Pre-image sampling (PIL `BILINEAR` / align-corners=False semantics): the diff --git a/swift/Sources/CoreAIShared/Runtime/VectorMath.swift b/swift/Sources/CoreAIShared/Runtime/VectorMath.swift new file mode 100644 index 00000000..18b40110 --- /dev/null +++ b/swift/Sources/CoreAIShared/Runtime/VectorMath.swift @@ -0,0 +1,19 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Foundation + +/// Numerically stable logistic sigmoid. +/// +/// Branches on the sign to keep the exponent argument non-positive, so unbounded +/// detector/segmenter logits stay finite at both extremes (no `NaN`, no underflow to `0`). +@inline(__always) +public func sigmoid(_ x: Float) -> Float { + if x >= 0 { + return 1 / (1 + Foundation.exp(-x)) + } + let e = Foundation.exp(x) + return e / (1 + e) +} diff --git a/swift/Sources/CoreAIVideoSegmenter/Tracking/DetectionDecoder.swift b/swift/Sources/CoreAIVideoSegmenter/Tracking/DetectionDecoder.swift index d7b1d011..a01812a0 100644 --- a/swift/Sources/CoreAIVideoSegmenter/Tracking/DetectionDecoder.swift +++ b/swift/Sources/CoreAIVideoSegmenter/Tracking/DetectionDecoder.swift @@ -109,15 +109,4 @@ enum DetectionDecoder { // prompt's own ordering. return kept.sorted() } - - @inline(__always) - static func sigmoid(_ x: Float) -> Float { - // Branch on the sign to keep the exponent argument negative. `exp` of a large - // positive logit overflows to infinity and yields NaN. - if x >= 0 { - return 1 / (1 + Foundation.exp(-x)) - } - let e = Foundation.exp(x) - return e / (1 + e) - } } diff --git a/swift/Sources/CoreAIVideoSegmenter/Tracking/FrameProcessor.swift b/swift/Sources/CoreAIVideoSegmenter/Tracking/FrameProcessor.swift index 72cf5a3f..085095dd 100644 --- a/swift/Sources/CoreAIVideoSegmenter/Tracking/FrameProcessor.swift +++ b/swift/Sources/CoreAIVideoSegmenter/Tracking/FrameProcessor.swift @@ -416,7 +416,7 @@ final class FrameProcessor { for (objectID, score) in newScores { trackerScores[objectID] = score } for (objectID, logit) in zip(trackedObjectIDs, trackerScoreLogits) where !removed.contains(objectID) { - trackerScores[objectID] = DetectionDecoder.sigmoid(logit) + trackerScores[objectID] = sigmoid(logit) } session.trackerScoreByFrame[frameIndex] = trackerScores diff --git a/swift/Tests/CoreAISharedTests/VectorMathTests.swift b/swift/Tests/CoreAISharedTests/VectorMathTests.swift new file mode 100644 index 00000000..eebcbfae --- /dev/null +++ b/swift/Tests/CoreAISharedTests/VectorMathTests.swift @@ -0,0 +1,37 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import Foundation +import Testing + +@testable import CoreAIShared + +// MARK: - sigmoid + +/// The shared logistic sigmoid is the single home for the two segmenter copies that previously +/// lived in `DetectionDecoder` and `SegmentationPostprocessor`. Detector/segmenter logits are +/// unbounded, so these cases pin the overflow-safe behavior at both extremes. +@Suite("VectorMath") +struct VectorMathTests { + @Test("sigmoid(0) == 0.5 and is symmetric") + func knownValues() { + #expect(sigmoid(0) == 0.5) + let x: Float = 2.5 + #expect(abs(sigmoid(x) + sigmoid(-x) - 1.0) < 1e-6) + #expect(abs(sigmoid(2) - 0.880797) < 1e-5) + } + + @Test("sigmoid stays finite at both extremes") + func overflowSafe() { + // Large positive: naive `exp(x)/(1+exp(x))` would be inf/inf = NaN; this returns 1. + #expect(sigmoid(100) == 1) + #expect(sigmoid(100).isNaN == false) + // Large negative: a tiny positive value, never NaN. The naive `1/(1+exp(-x))` underflows + // to exactly 0 here; the sign-branched form keeps the denormal. + #expect(sigmoid(-100) > 0) + #expect(sigmoid(-100) < 1e-40) + #expect(sigmoid(-100).isNaN == false) + } +} diff --git a/swift/Tests/ImageSegmenterTests/ImageSegmentationTests.swift b/swift/Tests/ImageSegmenterTests/ImageSegmentationTests.swift index dc473d1b..e7dd85d5 100644 --- a/swift/Tests/ImageSegmenterTests/ImageSegmentationTests.swift +++ b/swift/Tests/ImageSegmenterTests/ImageSegmentationTests.swift @@ -3,6 +3,7 @@ // Use of this source code is governed by a BSD-3-clause license that can // be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause +import CoreAIShared import CoreGraphics import Foundation import Testing @@ -95,7 +96,7 @@ struct ImageSegmentationTests { @Test("Decode: absent presence logits treated as 1.0") func decodeEmptyPresenceLogits() { let logit: Float = 2.0 - let expectedScore = SegmentationPostprocessor.sigmoid(logit) + let expectedScore = sigmoid(logit) let output = SegmentationOutput( predictedMasks: [Float](repeating: 0, count: 1 * 1 * 2 * 2), @@ -372,15 +373,6 @@ struct ImageSegmentationTests { #expect(overlay.height == h) } - // MARK: - Sigmoid - - @Test("sigmoid(0) == 0.5 and sigmoid(x) + sigmoid(-x) == 1") - func sigmoidKnownValues() { - #expect(SegmentationPostprocessor.sigmoid(0) == 0.5) - let x: Float = 2.5 - #expect(abs(SegmentationPostprocessor.sigmoid(x) + SegmentationPostprocessor.sigmoid(-x) - 1.0) < 1e-6) - } - // MARK: - Helpers private func makeSolidCGImage(width: Int, height: Int, r: UInt8, g: UInt8, b: UInt8) -> CGImage? { diff --git a/swift/Tests/VideoSegmenterTests/HotstartAndMemoryTests.swift b/swift/Tests/VideoSegmenterTests/HotstartAndMemoryTests.swift index aa8fba6a..9ce5e04f 100644 --- a/swift/Tests/VideoSegmenterTests/HotstartAndMemoryTests.swift +++ b/swift/Tests/VideoSegmenterTests/HotstartAndMemoryTests.swift @@ -588,9 +588,9 @@ struct BuildOutputsPairingTests { trackerLogits: maskLogits(for: tracked), trackerScoreLogits: scoreLogits(for: tracked), plan: plan, newScores: [:]) - #expect(output.trackerScoreByObjectID[1] == DetectionDecoder.sigmoid(1)) + #expect(output.trackerScoreByObjectID[1] == sigmoid(1)) #expect( - output.trackerScoreByObjectID[3] == DetectionDecoder.sigmoid(3), + output.trackerScoreByObjectID[3] == sigmoid(3), "object 3 must not inherit object 2's score") } diff --git a/swift/Tests/VideoSegmenterTests/TrackingPrimitiveTests.swift b/swift/Tests/VideoSegmenterTests/TrackingPrimitiveTests.swift index 9274bbf5..aeb81b5a 100644 --- a/swift/Tests/VideoSegmenterTests/TrackingPrimitiveTests.swift +++ b/swift/Tests/VideoSegmenterTests/TrackingPrimitiveTests.swift @@ -148,19 +148,6 @@ struct DetectionDecoderTests { return bitset } - @Test("sigmoid does not overflow on a large positive logit") - func sigmoidStability() { - // `1 / (1 + exp(-x))` is fine for large x but `exp(x) / (1 + exp(x))` overflows to - // inf/inf = NaN. The detector's logits are unbounded, so the branch matters. - #expect(DetectionDecoder.sigmoid(0) == 0.5) - #expect(DetectionDecoder.sigmoid(100) == 1) - // Not exactly zero: the negative branch keeps the denormal rather than underflowing. - #expect(DetectionDecoder.sigmoid(-100) < 1e-40) - #expect(DetectionDecoder.sigmoid(-100) >= 0) - #expect(DetectionDecoder.sigmoid(-100).isNaN == false) - #expect(abs(DetectionDecoder.sigmoid(2) - 0.880797) < 1e-5) - } - @Test("NMS keeps the highest-scoring member of an overlapping group") func nmsSuppresses() { let masks = [mask(["####"]), mask(["###."]), mask(["...#"])] @@ -235,7 +222,7 @@ struct DetectionDecoderTests { outputs(queries: [2], presence: 1, masks: [[1, 1, -1, -1]]), promptID: 0, maskSize: 2, parameters: parameters(threshold: 0.1, nms: 0)) #expect(decoded.count == 1) - let expected = DetectionDecoder.sigmoid(2) * DetectionDecoder.sigmoid(1) + let expected = sigmoid(2) * sigmoid(1) #expect(abs(decoded.scores[0] - expected) < 1e-6) #expect(abs(expected - 0.643914) < 1e-5) } @@ -289,8 +276,8 @@ struct DetectionDecoderTests { #expect(decoded.maskLogits[1] == [-1, -1, -1, 1]) #expect(decoded.masks[1].area == 1) let expected = [ - DetectionDecoder.sigmoid(3) * DetectionDecoder.sigmoid(8), - DetectionDecoder.sigmoid(1) * DetectionDecoder.sigmoid(8), + sigmoid(3) * sigmoid(8), + sigmoid(1) * sigmoid(8), ] #expect(zip(decoded.scores, expected).allSatisfy { abs($0 - $1) < 1e-6 }) }