mlx_c_rust/compile/
elementwise_graphs.rs

1use std::os::raw::c_int;
2
3use crate::MlxBindingsContext;
4use crate::compile::attention_output_gate::build_attention_output_gate_graph;
5use crate::compile::signed_float32_cast::build_signed_float32_cast_graph;
6use crate::compile::sparse_shared_expert_combination::build_sparse_shared_expert_combination_graph;
7use crate::compile::vision_rope::build_vision_rope_graph;
8use crate::{MlxArray, MlxStream, raw};
9use crate::{
10    MlxCError,
11    compile::graph::{MlxCompiledGraph, array_from_vector, graph_output_array, set_graph_output},
12};
13
14const APPLY_ATTENTION_OUTPUT_GATE_OPERATION: &str = "apply the compiled MLX attention output gate";
15const COMPILE_ATTENTION_OUTPUT_GATE_OPERATION: &str =
16    "compile the shapeless MLX attention output gate";
17const APPLY_SPARSE_SHARED_EXPERT_COMBINATION_OPERATION: &str =
18    "apply the compiled MLX sparse and shared expert combination";
19const COMPILE_SPARSE_SHARED_EXPERT_COMBINATION_OPERATION: &str =
20    "compile the shapeless MLX sparse and shared expert combination";
21const APPLY_PRECISE_SWIGLU_OPERATION: &str = "apply the compiled precise MLX SwiGLU graph";
22const COMPILE_PRECISE_SWIGLU_OPERATION: &str = "compile the shapeless precise MLX SwiGLU graph";
23const APPLY_GATED_DELTA_DECAY_OPERATION: &str =
24    "apply the compiled MLX gated-delta decay arithmetic";
25const COMPILE_GATED_DELTA_DECAY_OPERATION: &str =
26    "compile the shapeless MLX gated-delta decay arithmetic";
27const APPLY_FUSED_SILU_OPERATION: &str = "apply the compiled MLX fused SiLU graph";
28const COMPILE_FUSED_SILU_OPERATION: &str = "compile the shapeless MLX fused SiLU graph";
29const APPLY_FUSED_SOFTPLUS_OPERATION: &str = "apply the compiled MLX fused K2 softplus graph";
30const COMPILE_FUSED_SOFTPLUS_OPERATION: &str = "compile the shapeless MLX fused K2 softplus graph";
31const APPLY_SIGNED_FLOAT32_CAST_OPERATION: &str = "apply the compiled MLX signed float32 cast";
32const COMPILE_SIGNED_FLOAT32_CAST_OPERATION: &str = "compile the shapeless MLX signed float32 cast";
33const APPLY_VISION_ROPE_OPERATION: &str =
34    "apply the compiled MLX vision rotate-half rotary embedding";
35const COMPILE_VISION_ROPE_OPERATION: &str =
36    "compile the shapeless MLX vision rotate-half rotary embedding";
37
38/// Retained shapeless compilations for elementwise model graph composites.
39#[derive(Debug)]
40pub struct MlxCompiledElementwiseGraphs {
41    attention_output_gate: MlxCompiledGraph,
42    gated_delta_decay: MlxCompiledGraph,
43    precise_swiglu: MlxCompiledGraph,
44    sparse_shared_expert_combination: MlxCompiledGraph,
45    fused_silu: MlxCompiledGraph,
46    fused_softplus: MlxCompiledGraph,
47    vision_rope: MlxCompiledGraph,
48    signed_float32_cast: MlxCompiledGraph,
49}
50
51impl MlxCompiledElementwiseGraphs {
52    /// Creates reusable compiled graphs for elementwise model composites.
53    pub fn new() -> Result<Self, MlxCError> {
54        Ok(Self {
55            attention_output_gate: MlxCompiledGraph::new(
56                build_attention_output_gate_graph,
57                COMPILE_ATTENTION_OUTPUT_GATE_OPERATION,
58                true,
59            )?,
60            gated_delta_decay: MlxCompiledGraph::new(
61                build_gated_delta_decay_graph,
62                COMPILE_GATED_DELTA_DECAY_OPERATION,
63                true,
64            )?,
65            precise_swiglu: MlxCompiledGraph::new(
66                build_precise_swiglu_graph,
67                COMPILE_PRECISE_SWIGLU_OPERATION,
68                true,
69            )?,
70            sparse_shared_expert_combination: MlxCompiledGraph::new(
71                build_sparse_shared_expert_combination_graph,
72                COMPILE_SPARSE_SHARED_EXPERT_COMBINATION_OPERATION,
73                true,
74            )?,
75            fused_silu: MlxCompiledGraph::new(
76                build_fused_silu_graph,
77                COMPILE_FUSED_SILU_OPERATION,
78                true,
79            )?,
80            fused_softplus: MlxCompiledGraph::new(
81                build_fused_softplus_graph,
82                COMPILE_FUSED_SOFTPLUS_OPERATION,
83                true,
84            )?,
85            vision_rope: MlxCompiledGraph::new(
86                build_vision_rope_graph,
87                COMPILE_VISION_ROPE_OPERATION,
88                true,
89            )?,
90            signed_float32_cast: MlxCompiledGraph::new(
91                build_signed_float32_cast_graph,
92                COMPILE_SIGNED_FLOAT32_CAST_OPERATION,
93                true,
94            )?,
95        })
96    }
97}
98
99impl MlxBindingsContext {
100    /// Applies `attention_output * sigmoid(output_gate_logits)` as one compiled composite.
101    pub fn apply_compiled_attention_output_gate(
102        &self,
103        compiled_elementwise_graphs: &MlxCompiledElementwiseGraphs,
104        attention_output: &MlxArray,
105        output_gate_logits: &MlxArray,
106    ) -> Result<MlxArray, MlxCError> {
107        compiled_elementwise_graphs.attention_output_gate.apply(
108            &[attention_output, output_gate_logits],
109            APPLY_ATTENTION_OUTPUT_GATE_OPERATION,
110        )
111    }
112
113    /// Applies `sparse + shared * sigmoid(shared_gate_logits)` as one compiled composite.
114    pub fn apply_compiled_sparse_shared_expert_combination(
115        &self,
116        compiled_elementwise_graphs: &MlxCompiledElementwiseGraphs,
117        sparse_expert_output: &MlxArray,
118        shared_expert_output: &MlxArray,
119        shared_expert_gate_logits: &MlxArray,
120    ) -> Result<MlxArray, MlxCError> {
121        compiled_elementwise_graphs
122            .sparse_shared_expert_combination
123            .apply(
124                &[
125                    sparse_expert_output,
126                    shared_expert_output,
127                    shared_expert_gate_logits,
128                ],
129                APPLY_SPARSE_SHARED_EXPERT_COMBINATION_OPERATION,
130            )
131    }
132
133    /// Applies float32 SwiGLU arithmetic and restores the input activation dtype.
134    pub fn apply_compiled_precise_swiglu(
135        &self,
136        compiled_elementwise_graphs: &MlxCompiledElementwiseGraphs,
137        up_states: &MlxArray,
138        gate_states: &MlxArray,
139    ) -> Result<MlxArray, MlxCError> {
140        compiled_elementwise_graphs
141            .precise_swiglu
142            .apply(&[up_states, gate_states], APPLY_PRECISE_SWIGLU_OPERATION)
143    }
144
145    /// Applies the complete float32 gated-delta decay formula as one compiled composite.
146    pub fn apply_compiled_gated_delta_decay(
147        &self,
148        compiled_elementwise_graphs: &MlxCompiledElementwiseGraphs,
149        decay_rate_logarithm: &MlxArray,
150        decay_interval_inputs: &MlxArray,
151        decay_interval_bias: &MlxArray,
152    ) -> Result<MlxArray, MlxCError> {
153        compiled_elementwise_graphs.gated_delta_decay.apply(
154            &[
155                decay_rate_logarithm,
156                decay_interval_inputs,
157                decay_interval_bias,
158            ],
159            APPLY_GATED_DELTA_DECAY_OPERATION,
160        )
161    }
162
163    /// Applies `input * sigmoid(input)` as one compiled composite.
164    ///
165    /// Stock MLX defines `nn.silu` as a shapeless `mx.compile`, which fuses the
166    /// sigmoid and the multiply into one kernel reading the input once. The raw
167    /// two-op composition runs two kernels and roughly 2.5x the memory traffic,
168    /// which the gated-delta convolution path paid on every layer of every
169    /// prefill chunk and decode step.
170    pub fn apply_compiled_silu(
171        &self,
172        compiled_elementwise_graphs: &MlxCompiledElementwiseGraphs,
173        input: &MlxArray,
174    ) -> Result<MlxArray, MlxCError> {
175        compiled_elementwise_graphs
176            .fused_silu
177            .apply(&[input], APPLY_FUSED_SILU_OPERATION)
178    }
179
180    /// Applies `astype(activations, float32) * signs` as one compiled composite.
181    ///
182    /// The Hadamard-rotated families cast every activation to float32 and
183    /// flip its signs before the block transform; uncompiled, that is a cast
184    /// kernel plus a multiply kernel with a float32 intermediate round trip.
185    /// The fused launch reads the activation once and writes the signed
186    /// float32 result once, with arithmetic identical to the two-op chain.
187    pub fn apply_compiled_signed_float32_cast(
188        &self,
189        compiled_elementwise_graphs: &MlxCompiledElementwiseGraphs,
190        activations: &MlxArray,
191        signs: &MlxArray,
192    ) -> Result<MlxArray, MlxCError> {
193        compiled_elementwise_graphs
194            .signed_float32_cast
195            .apply(&[activations, signs], APPLY_SIGNED_FLOAT32_CAST_OPERATION)
196    }
197
198    /// Applies the K2 Horizon MoVA softplus gate as one compiled composite.
199    ///
200    /// The gate formula is `log1p(exp(input * ln2)) / ln2` — the reference
201    /// implementation's beta-scaled softplus, which is NOT the standard
202    /// `logaddexp(input, 0)` softplus and must not be "simplified" to it.
203    /// Compiling collapses the four hand-composed dispatches into one launch
204    /// without changing the arithmetic.
205    pub fn apply_compiled_softplus(
206        &self,
207        compiled_elementwise_graphs: &MlxCompiledElementwiseGraphs,
208        input: &MlxArray,
209    ) -> Result<MlxArray, MlxCError> {
210        compiled_elementwise_graphs
211            .fused_softplus
212            .apply(&[input], APPLY_FUSED_SOFTPLUS_OPERATION)
213    }
214
215    /// Applies Qwen3-VL rotate-half rotary embedding as one compiled composite.
216    ///
217    /// `first_half` and `second_half` are the pre-sliced head halves of
218    /// `attention_states` (the last axis split at `head_dim / 2`); the head-half
219    /// slices stay outside the compiled graph because the pinned MLX `Slice`
220    /// primitive cannot be shapeless-compiled. The remaining composite fuses the
221    /// elementwise tail into one kernel and keeps the rotate-half concat as one
222    /// shape kernel, replacing the eight separate dispatches.
223    pub fn apply_compiled_vision_rope(
224        &self,
225        compiled_elementwise_graphs: &MlxCompiledElementwiseGraphs,
226        attention_states: &MlxArray,
227        rotary_cosines: &MlxArray,
228        rotary_sines: &MlxArray,
229        first_half: &MlxArray,
230        second_half: &MlxArray,
231    ) -> Result<MlxArray, MlxCError> {
232        compiled_elementwise_graphs.vision_rope.apply(
233            &[
234                attention_states,
235                rotary_cosines,
236                rotary_sines,
237                first_half,
238                second_half,
239            ],
240            APPLY_VISION_ROPE_OPERATION,
241        )
242    }
243}
244
245unsafe extern "C" fn build_precise_swiglu_graph(
246    output_vector: *mut raw::mlx_vector_array,
247    input_vector: raw::mlx_vector_array,
248) -> c_int {
249    if output_vector.is_null() || unsafe { raw::mlx_vector_array_size(input_vector) } != 2 {
250        return 1;
251    }
252    let up_states = match array_from_vector(input_vector, 0) {
253        Ok(up_states) => up_states,
254        Err(get_status) => return get_status,
255    };
256    let gate_states = match array_from_vector(input_vector, 1) {
257        Ok(gate_states) => gate_states,
258        Err(get_status) => return get_status,
259    };
260    let gpu_stream = match MlxStream::default_gpu() {
261        Ok(gpu_stream) => gpu_stream,
262        Err(_) => return 1,
263    };
264    let output_dtype = unsafe { raw::mlx_array_dtype(up_states.raw()) };
265    let float32_gate = match graph_output_array(|output_array| {
266        // SAFETY: The input and stream are live, and the output is uniquely writable.
267        unsafe {
268            raw::mlx_astype(
269                output_array,
270                gate_states.raw(),
271                raw::mlx_dtype__MLX_FLOAT32,
272                gpu_stream.raw(),
273            )
274        }
275    }) {
276        Ok(float32_gate) => float32_gate,
277        Err(build_status) => return build_status,
278    };
279    let gate_weights = match graph_output_array(|output_array| {
280        // SAFETY: The input and stream are live, and the output is uniquely writable.
281        unsafe { raw::mlx_sigmoid(output_array, float32_gate.raw(), gpu_stream.raw()) }
282    }) {
283        Ok(gate_weights) => gate_weights,
284        Err(build_status) => return build_status,
285    };
286    let activated_gate = match graph_output_array(|output_array| {
287        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
288        unsafe {
289            raw::mlx_multiply(
290                output_array,
291                float32_gate.raw(),
292                gate_weights.raw(),
293                gpu_stream.raw(),
294            )
295        }
296    }) {
297        Ok(activated_gate) => activated_gate,
298        Err(build_status) => return build_status,
299    };
300    let float32_up = match graph_output_array(|output_array| {
301        // SAFETY: The input and stream are live, and the output is uniquely writable.
302        unsafe {
303            raw::mlx_astype(
304                output_array,
305                up_states.raw(),
306                raw::mlx_dtype__MLX_FLOAT32,
307                gpu_stream.raw(),
308            )
309        }
310    }) {
311        Ok(float32_up) => float32_up,
312        Err(build_status) => return build_status,
313    };
314    let float32_activated_states = match graph_output_array(|output_array| {
315        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
316        unsafe {
317            raw::mlx_multiply(
318                output_array,
319                activated_gate.raw(),
320                float32_up.raw(),
321                gpu_stream.raw(),
322            )
323        }
324    }) {
325        Ok(float32_activated_states) => float32_activated_states,
326        Err(build_status) => return build_status,
327    };
328    let activated_states = match graph_output_array(|output_array| {
329        // SAFETY: The input and stream are live, and the output is uniquely writable.
330        unsafe {
331            raw::mlx_astype(
332                output_array,
333                float32_activated_states.raw(),
334                output_dtype,
335                gpu_stream.raw(),
336            )
337        }
338    }) {
339        Ok(activated_states) => activated_states,
340        Err(build_status) => return build_status,
341    };
342    // SAFETY: The output vector is unique and live for this callback.
343    unsafe { set_graph_output(output_vector, &activated_states) }
344}
345
346unsafe extern "C" fn build_gated_delta_decay_graph(
347    output_vector: *mut raw::mlx_vector_array,
348    input_vector: raw::mlx_vector_array,
349) -> c_int {
350    if output_vector.is_null() || unsafe { raw::mlx_vector_array_size(input_vector) } != 3 {
351        return 1;
352    }
353    let decay_rate_logarithm = match array_from_vector(input_vector, 0) {
354        Ok(decay_rate_logarithm) => decay_rate_logarithm,
355        Err(get_status) => return get_status,
356    };
357    let decay_interval_inputs = match array_from_vector(input_vector, 1) {
358        Ok(decay_interval_inputs) => decay_interval_inputs,
359        Err(get_status) => return get_status,
360    };
361    let decay_interval_bias = match array_from_vector(input_vector, 2) {
362        Ok(decay_interval_bias) => decay_interval_bias,
363        Err(get_status) => return get_status,
364    };
365    let gpu_stream = match MlxStream::default_gpu() {
366        Ok(gpu_stream) => gpu_stream,
367        Err(_) => return 1,
368    };
369    let biased_decay_intervals = match graph_output_array(|output_array| {
370        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
371        unsafe {
372            raw::mlx_add(
373                output_array,
374                decay_interval_inputs.raw(),
375                decay_interval_bias.raw(),
376                gpu_stream.raw(),
377            )
378        }
379    }) {
380        Ok(biased_decay_intervals) => biased_decay_intervals,
381        Err(build_status) => return build_status,
382    };
383    // Stable softplus through MLX logaddexp(x, 0).
384    // The naive log1p(exp(x)) overflows to infinity for large positive decay
385    // intervals, corrupting the gated-delta recurrent state during long-context
386    // prefill. This matches MLX's nn.softplus, which delegates to logaddexp(x, 0).
387    let biased_dtype = biased_decay_intervals.dtype();
388    let scalar_shape_storage = [1_i32];
389    let zero_decay_interval_scalar = match graph_output_array(|output_array| {
390        // SAFETY: The shape pointer references live storage but the zero rank
391        // means MLX reads no dimensions. Dtype and stream are valid, and the
392        // output is uniquely writable. Broadcasting this scalar keeps the
393        // compiled graph shapeless across variable token counts.
394        unsafe {
395            raw::mlx_zeros(
396                output_array,
397                scalar_shape_storage.as_ptr(),
398                0,
399                biased_dtype.to_raw(),
400                gpu_stream.raw(),
401            )
402        }
403    }) {
404        Ok(zero_decay_interval_scalar) => zero_decay_interval_scalar,
405        Err(build_status) => return build_status,
406    };
407    let decay_intervals = match graph_output_array(|output_array| {
408        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
409        unsafe {
410            raw::mlx_logaddexp(
411                output_array,
412                biased_decay_intervals.raw(),
413                zero_decay_interval_scalar.raw(),
414                gpu_stream.raw(),
415            )
416        }
417    }) {
418        Ok(decay_intervals) => decay_intervals,
419        Err(build_status) => return build_status,
420    };
421    let float32_decay_logs = match graph_output_array(|output_array| {
422        // SAFETY: The input and stream are live, and the output is uniquely writable.
423        unsafe {
424            raw::mlx_astype(
425                output_array,
426                decay_rate_logarithm.raw(),
427                raw::mlx_dtype__MLX_FLOAT32,
428                gpu_stream.raw(),
429            )
430        }
431    }) {
432        Ok(float32_decay_logs) => float32_decay_logs,
433        Err(build_status) => return build_status,
434    };
435    let decay_rates = match graph_output_array(|output_array| {
436        // SAFETY: The input and stream are live, and the output is uniquely writable.
437        unsafe { raw::mlx_exp(output_array, float32_decay_logs.raw(), gpu_stream.raw()) }
438    }) {
439        Ok(decay_rates) => decay_rates,
440        Err(build_status) => return build_status,
441    };
442    let decay_products = match graph_output_array(|output_array| {
443        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
444        unsafe {
445            raw::mlx_multiply(
446                output_array,
447                decay_rates.raw(),
448                decay_intervals.raw(),
449                gpu_stream.raw(),
450            )
451        }
452    }) {
453        Ok(decay_products) => decay_products,
454        Err(build_status) => return build_status,
455    };
456    let negative_decay_products = match graph_output_array(|output_array| {
457        // SAFETY: The input and stream are live, and the output is uniquely writable.
458        unsafe { raw::mlx_negative(output_array, decay_products.raw(), gpu_stream.raw()) }
459    }) {
460        Ok(negative_decay_products) => negative_decay_products,
461        Err(build_status) => return build_status,
462    };
463    let decays = match graph_output_array(|output_array| {
464        // SAFETY: The input and stream are live, and the output is uniquely writable.
465        unsafe {
466            raw::mlx_exp(
467                output_array,
468                negative_decay_products.raw(),
469                gpu_stream.raw(),
470            )
471        }
472    }) {
473        Ok(decays) => decays,
474        Err(build_status) => return build_status,
475    };
476    // SAFETY: The output vector is unique and live for this callback.
477    unsafe { set_graph_output(output_vector, &decays) }
478}
479
480unsafe extern "C" fn build_fused_silu_graph(
481    output_vector: *mut raw::mlx_vector_array,
482    input_vector: raw::mlx_vector_array,
483) -> c_int {
484    if output_vector.is_null() || unsafe { raw::mlx_vector_array_size(input_vector) } != 1 {
485        return 1;
486    }
487    let input = match array_from_vector(input_vector, 0) {
488        Ok(input) => input,
489        Err(get_status) => return get_status,
490    };
491    let gpu_stream = match MlxStream::default_gpu() {
492        Ok(gpu_stream) => gpu_stream,
493        Err(_) => return 1,
494    };
495    let sigmoid_input = match graph_output_array(|output_array| {
496        // SAFETY: The input and stream are live, and the output is uniquely writable.
497        unsafe { raw::mlx_sigmoid(output_array, input.raw(), gpu_stream.raw()) }
498    }) {
499        Ok(sigmoid_input) => sigmoid_input,
500        Err(build_status) => return build_status,
501    };
502    let silu_output = match graph_output_array(|output_array| {
503        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
504        unsafe {
505            raw::mlx_multiply(
506                output_array,
507                input.raw(),
508                sigmoid_input.raw(),
509                gpu_stream.raw(),
510            )
511        }
512    }) {
513        Ok(silu_output) => silu_output,
514        Err(build_status) => return build_status,
515    };
516    // SAFETY: The output vector is unique and live for this callback.
517    unsafe { set_graph_output(output_vector, &silu_output) }
518}
519
520unsafe extern "C" fn build_fused_softplus_graph(
521    output_vector: *mut raw::mlx_vector_array,
522    input_vector: raw::mlx_vector_array,
523) -> c_int {
524    if output_vector.is_null() || unsafe { raw::mlx_vector_array_size(input_vector) } != 1 {
525        return 1;
526    }
527    let input = match array_from_vector(input_vector, 0) {
528        Ok(input) => input,
529        Err(get_status) => return get_status,
530    };
531    let gpu_stream = match MlxStream::default_gpu() {
532        Ok(gpu_stream) => gpu_stream,
533        Err(_) => return 1,
534    };
535    let ln2 = std::f32::consts::LN_2;
536    let ln2_scalar = match MlxArray::from_f32(&[ln2], &[]) {
537        Ok(ln2_scalar) => ln2_scalar,
538        Err(_) => return 1,
539    };
540    let typed_ln2 = match graph_output_array(|output_array| {
541        // SAFETY: The scalar and stream are live, and the output is uniquely writable.
542        unsafe {
543            raw::mlx_astype(
544                output_array,
545                ln2_scalar.raw(),
546                raw::mlx_array_dtype(input.raw()),
547                gpu_stream.raw(),
548            )
549        }
550    }) {
551        Ok(typed_ln2) => typed_ln2,
552        Err(build_status) => return build_status,
553    };
554    let scaled_input = match graph_output_array(|output_array| {
555        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
556        unsafe { raw::mlx_multiply(output_array, input.raw(), typed_ln2.raw(), gpu_stream.raw()) }
557    }) {
558        Ok(scaled_input) => scaled_input,
559        Err(build_status) => return build_status,
560    };
561    let exponent = match graph_output_array(|output_array| {
562        // SAFETY: The input and stream are live, and the output is uniquely writable.
563        unsafe { raw::mlx_exp(output_array, scaled_input.raw(), gpu_stream.raw()) }
564    }) {
565        Ok(exponent) => exponent,
566        Err(build_status) => return build_status,
567    };
568    let logarithm = match graph_output_array(|output_array| {
569        // SAFETY: The input and stream are live, and the output is uniquely writable.
570        unsafe { raw::mlx_log1p(output_array, exponent.raw(), gpu_stream.raw()) }
571    }) {
572        Ok(logarithm) => logarithm,
573        Err(build_status) => return build_status,
574    };
575    let inverse_ln2_scalar = match MlxArray::from_f32(&[1.0 / ln2], &[]) {
576        Ok(inverse_ln2_scalar) => inverse_ln2_scalar,
577        Err(_) => return 1,
578    };
579    let typed_inverse_ln2 = match graph_output_array(|output_array| {
580        // SAFETY: The scalar and stream are live, and the output is uniquely writable.
581        unsafe {
582            raw::mlx_astype(
583                output_array,
584                inverse_ln2_scalar.raw(),
585                raw::mlx_array_dtype(input.raw()),
586                gpu_stream.raw(),
587            )
588        }
589    }) {
590        Ok(typed_inverse_ln2) => typed_inverse_ln2,
591        Err(build_status) => return build_status,
592    };
593    let softplus_output = match graph_output_array(|output_array| {
594        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
595        unsafe {
596            raw::mlx_multiply(
597                output_array,
598                logarithm.raw(),
599                typed_inverse_ln2.raw(),
600                gpu_stream.raw(),
601            )
602        }
603    }) {
604        Ok(softplus_output) => softplus_output,
605        Err(build_status) => return build_status,
606    };
607    // SAFETY: The output vector is unique and live for this callback.
608    unsafe { set_graph_output(output_vector, &softplus_output) }
609}
mlx_c_rust/compile/elementwise_graphs.rs

mlx_c_rust/compile/
elementwise_graphs.rs
