+3

Libraries/MLXLLM/Models/BailingMoe.swift

View file Reviewed

··· 130 130 let kSize = kvHeads * headDim 131 131 let qkvOut = qkv(x) 132 132 let splits = split(qkvOut, indices: [qSize, qSize + kSize], axis: -1) 133 133 + // Bail on a failed split (empty vector from a recorded MLX error) 134 134 + // rather than trapping on the q/k/v subscripts below. 135 135 + guard splits.count == 3 else { return x } 133 136 var queries = splits[0] 134 137 var keys = splits[1] 135 138 var values = splits[2]

+3

Libraries/MLXLLM/Models/FalconH1.swift

View file Reviewed

··· 474 474 indices: [intermediateSize, intermediateSize + convDim], 475 475 axis: -1 476 476 ) 477 477 + // Bail on a failed split rather than trapping on the subscripts below. 478 478 + guard splits.count == 3 else { return inputStates } 477 479 let gate = splits[0] 478 480 var convInput = splits[1] 479 481 let dt = splits[2] ··· 491 493 ], 492 494 axis: -1 493 495 ) 496 496 + guard convSplits.count == 3 else { return inputStates } 494 497 let hiddenStatesSSM = convSplits[0] 495 498 let B = convSplits[1] 496 499 let C = convSplits[2]

+3

Libraries/MLXLLM/Models/GraniteMoeHybrid.swift

View file Reviewed

··· 139 139 let projected = inProj(hiddenStates) 140 140 let splits = split( 141 141 projected, indices: [intermediateSize, intermediateSize + convDim], axis: -1) 142 142 + // Bail on a failed split rather than trapping on the subscripts below. 143 143 + guard splits.count == 3 else { return hiddenStates } 142 144 let gate = splits[0] 143 145 var convInput = splits[1] 144 146 let dt = splits[2] ··· 155 157 axis: -1 156 158 ) 157 159 160 160 + guard convSplits.count == 3 else { return hiddenStates } 158 161 var hidden = convSplits[0] 159 162 var B = convSplits[1] 160 163 var C = convSplits[2]

+4

Libraries/MLXLLM/Models/Jamba.swift

View file Reviewed

··· 264 264 indices: [timeStepRank, timeStepRank + ssmStateSize], 265 265 axis: -1 266 266 ) 267 267 + // Bail on a failed split rather than trapping on the subscripts below. 268 268 + guard splits.count == 3 else { return (x, state ?? x) } 267 269 var delta = splits[0] 268 270 var B = splits[1] 269 271 var C = splits[2] ··· 294 296 ) { 295 297 let xz = inProj(x) 296 298 let splits = xz.split(parts: 2, axis: -1) 299 299 + // Bail on a failed split rather than trapping on the subscripts below. 300 300 + guard splits.count == 2 else { return (x, (x, x)) } 297 301 var x = splits[0] 298 302 let z = splits[1] 299 303

+3

Libraries/MLXLLM/Models/LFM2.swift

View file Reviewed

··· 206 206 public func callAsFunction(_ x: MLXArray, cache: MambaCache?) -> MLXArray { 207 207 let BCx = inProj(x) 208 208 let BCxSplit = BCx.split(parts: 3, axis: -1) 209 209 + // Bail on a failed split (empty vector from a recorded MLX error) 210 210 + // rather than trapping on the subscripts below. 211 211 + guard BCxSplit.count == 3 else { return x } 209 212 let B = BCxSplit[0] 210 213 let C = BCxSplit[1] 211 214 let x = BCxSplit[2]

+2

Libraries/MLXLLM/Models/LFM2MoE.swift

View file Reviewed

··· 207 207 ) -> MLXArray { 208 208 let BCx = inProj(x) 209 209 let parts = BCx.split(parts: 3, axis: -1) 210 210 + // Bail on a failed split rather than trapping on the subscripts below. 211 211 + guard parts.count == 3 else { return x } 210 212 let B = parts[0] 211 213 let C = parts[1] 212 214 let xComp = parts[2]

+9

Libraries/MLXLLM/Models/Qwen35.swift

View file Reviewed

··· 290 290 let convOut = silu(conv1d(convInput)) 291 291 292 292 let convSplit = MLX.split(convOut, indices: [keyDim, 2 * keyDim], axis: -1) 293 293 + // A failed split (e.g. a checkpoint whose conv width disagrees with 294 294 + // keyDim*2 + valueDim) records an MLX error and returns an empty 295 295 + // vector; subscripting it traps the process before the recorded error 296 296 + // can surface. Bail with the input so the enclosing withError scope 297 297 + // throws the real diagnostic at exit. (Same guard as NemotronH.) 298 298 + guard convSplit.count == 3 else { return inputs } 293 299 let q = convSplit[0].reshaped(B, S, numKHeads, headKDim) 294 300 let k = convSplit[1].reshaped(B, S, numKHeads, headKDim) 295 301 let v = convSplit[2].reshaped(B, S, numVHeads, headVDim) ··· 465 471 466 472 let qProjOutput = qProj(x) 467 473 let qSplit = qProjOutput.reshaped(B, L, attentionHeads, -1).split(parts: 2, axis: -1) 474 474 + // Guard the query/gate split: an empty result from a recorded MLX error 475 475 + // would trap on subscript before the diagnostic can surface. 476 476 + guard qSplit.count == 2 else { return x } 468 477 var queries = qSplit[0] 469 478 let gate = qSplit[1].reshaped(B, L, -1) 470 479

+8

Libraries/MLXLLM/Models/Qwen3Next.swift

View file Reviewed

··· 121 121 122 122 let qProjOutput = qProj(x) 123 123 let qSplit = qProjOutput.reshaped(B, L, args.attentionHeads, -1).split(parts: 2, axis: -1) 124 124 + // Bail on a failed split rather than trapping on the subscripts below. 125 125 + guard qSplit.count == 2 else { return x } 124 126 var queries = qSplit[0] 125 127 let gate = qSplit[1].reshaped(B, L, -1) 126 128 ··· 246 248 indices: [dn, 2 * dn, 2 * dn + vHeadsPerK * dv], 247 249 axis: -1 248 250 ) 251 251 + // Bail on a failed split rather than trapping on the subscripts below; 252 252 + // the recorded MLX error surfaces at the next eval. 253 253 + guard qkvzSplit.count == 4 else { return (qkvz, qkvz, qkvz, qkvz, ba, ba) } 249 254 let q = qkvzSplit[0] 250 255 let k = qkvzSplit[1] 251 256 let v = qkvzSplit[2].reshaped(B, S, -1, dv) 252 257 let z = qkvzSplit[3].reshaped(B, S, -1, dv) 253 258 254 259 let baSplit = MLX.split(ba, indices: [vHeadsPerK], axis: -1) 260 260 + guard baSplit.count == 2 else { return (q, k, v, z, ba, ba) } 255 261 let b = baSplit[0].reshaped(B, S, nv) 256 262 let a = baSplit[1].reshaped(B, S, nv) 257 263 ··· 297 303 let convOut = silu(conv1d(convInput)) 298 304 let convSplit = MLX.split(convOut, indices: [keyDim, 2 * keyDim], axis: -1) 299 305 306 306 + // Bail on a failed split rather than trapping on the subscripts below. 307 307 + guard convSplit.count == 3 else { return inputs } 300 308 var qOut = convSplit[0].reshaped(B, S, numKHeads, headKDim) 301 309 var kOut = convSplit[1].reshaped(B, S, numKHeads, headKDim) 302 310 let vOut = convSplit[2].reshaped(B, S, numVHeads, headVDim)