Skip to content

Commit 2a34cde

Browse files
committed
Remove redundant broadcast op — MPSGraph elementwise ops broadcast implicitly
1 parent 67705e4 commit 2a34cde

1 file changed

Lines changed: 1 addition & 2 deletions

File tree

swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -620,8 +620,7 @@ final class MPSGraphCompositeSampler: @unchecked Sendable {
620620

621621
// Step 8: Multinomial sampling via cumulative sum + random comparison
622622
let cumsum = graph.cumulativeSum(normalizedProbs, axis: 1, exclusive: false, reverse: false, name: "cumsum")
623-
let randomBroadcast = graph.broadcast(randomPlaceholder, shape: [1, k as NSNumber], name: "random_broadcast")
624-
let selectionMask = graph.greaterThanOrEqualTo(cumsum, randomBroadcast, name: "selection_mask")
623+
let selectionMask = graph.greaterThanOrEqualTo(cumsum, randomPlaceholder, name: "selection_mask")
625624
let selectionMaskFloat = graph.cast(selectionMask, to: .float32, name: "selection_mask_float")
626625
let selectedIdx = graph.reductionArgMaximum(with: selectionMaskFloat, axis: 1, name: "selected_idx")
627626

0 commit comments

Comments
 (0)