Skip to main content

runmat_runtime/builtins/math/reduction/
gradient.rs

1//! MATLAB-compatible `gradient` builtin with scalar-spacing GPU residency.
2
3use runmat_accelerate_api::{GpuTensorHandle, GpuTensorStorage};
4use runmat_builtins::{
5    BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
6    BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
7    ComplexTensor, ResolveContext, Tensor, Type, Value,
8};
9use runmat_macros::runtime_builtin;
10
11use crate::builtins::common::gpu_helpers;
12use crate::builtins::common::random_args::complex_tensor_into_value;
13use crate::builtins::common::spec::{
14    BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
15    ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
16};
17use crate::builtins::common::tensor;
18use crate::builtins::math::type_resolvers::numeric_unary_type;
19use crate::{build_runtime_error, BuiltinResult, RuntimeError};
20
21const NAME: &str = "gradient";
22
23fn gradient_type(args: &[Type], ctx: &ResolveContext) -> Type {
24    numeric_unary_type(args, ctx)
25}
26
27const GRADIENT_OUTPUT_G: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
28    name: "G",
29    ty: BuiltinParamType::NumericArray,
30    arity: BuiltinParamArity::Required,
31    default: None,
32    description: "Primary gradient component.",
33}];
34
35const GRADIENT_OUTPUT_GS: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
36    name: "Gi",
37    ty: BuiltinParamType::NumericArray,
38    arity: BuiltinParamArity::Variadic,
39    default: None,
40    description: "Gradient components ordered by MATLAB axis semantics.",
41}];
42
43const GRADIENT_INPUTS_F: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
44    name: "F",
45    ty: BuiltinParamType::Any,
46    arity: BuiltinParamArity::Required,
47    default: None,
48    description: "Input scalar or array.",
49}];
50
51const GRADIENT_INPUTS_F_H: [BuiltinParamDescriptor; 2] = [
52    BuiltinParamDescriptor {
53        name: "F",
54        ty: BuiltinParamType::Any,
55        arity: BuiltinParamArity::Required,
56        default: None,
57        description: "Input scalar or array.",
58    },
59    BuiltinParamDescriptor {
60        name: "h",
61        ty: BuiltinParamType::Any,
62        arity: BuiltinParamArity::Optional,
63        default: Some("1"),
64        description: "Scalar spacing shared across all output dimensions.",
65    },
66];
67
68const GRADIENT_INPUTS_F_HS: [BuiltinParamDescriptor; 2] = [
69    BuiltinParamDescriptor {
70        name: "F",
71        ty: BuiltinParamType::Any,
72        arity: BuiltinParamArity::Required,
73        default: None,
74        description: "Input scalar or array.",
75    },
76    BuiltinParamDescriptor {
77        name: "h_i",
78        ty: BuiltinParamType::Any,
79        arity: BuiltinParamArity::Variadic,
80        default: None,
81        description: "Per-dimension scalar spacings (one per requested gradient component).",
82    },
83];
84
85const GRADIENT_SIGNATURES: [BuiltinSignatureDescriptor; 4] = [
86    BuiltinSignatureDescriptor {
87        label: "G = gradient(F)",
88        inputs: &GRADIENT_INPUTS_F,
89        outputs: &GRADIENT_OUTPUT_G,
90    },
91    BuiltinSignatureDescriptor {
92        label: "G = gradient(F, h)",
93        inputs: &GRADIENT_INPUTS_F_H,
94        outputs: &GRADIENT_OUTPUT_G,
95    },
96    BuiltinSignatureDescriptor {
97        label: "[G1, G2, ...] = gradient(F)",
98        inputs: &GRADIENT_INPUTS_F,
99        outputs: &GRADIENT_OUTPUT_GS,
100    },
101    BuiltinSignatureDescriptor {
102        label: "[G1, G2, ...] = gradient(F, h1, h2, ...)",
103        inputs: &GRADIENT_INPUTS_F_HS,
104        outputs: &GRADIENT_OUTPUT_GS,
105    },
106];
107
108const GRADIENT_ERROR_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
109    code: "RM.GRADIENT.INVALID_ARGUMENT",
110    identifier: Some("RunMat:gradient:InvalidArgument"),
111    when: "Output-count or spacing argument grammar is invalid.",
112    message: "gradient: invalid argument",
113};
114
115const GRADIENT_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
116    code: "RM.GRADIENT.INVALID_INPUT",
117    identifier: Some("RunMat:gradient:InvalidInput"),
118    when: "Input value cannot be converted to a supported gradient domain.",
119    message: "gradient: invalid input",
120};
121
122const GRADIENT_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
123    code: "RM.GRADIENT.INTERNAL",
124    identifier: Some("RunMat:gradient:Internal"),
125    when: "Gradient execution fails due to gather, conversion, allocation, or indexing operations.",
126    message: "gradient: internal failure",
127};
128
129const GRADIENT_ERRORS: [BuiltinErrorDescriptor; 3] = [
130    GRADIENT_ERROR_INVALID_ARGUMENT,
131    GRADIENT_ERROR_INVALID_INPUT,
132    GRADIENT_ERROR_INTERNAL,
133];
134
135pub const GRADIENT_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
136    signatures: &GRADIENT_SIGNATURES,
137    output_mode: BuiltinOutputMode::ByRequestedOutputCount,
138    completion_policy: BuiltinCompletionPolicy::Public,
139    errors: &GRADIENT_ERRORS,
140};
141
142fn gradient_descriptor_error_with_message(
143    message: impl Into<String>,
144    error: &'static BuiltinErrorDescriptor,
145) -> RuntimeError {
146    let mut builder = build_runtime_error(message).with_builtin(NAME);
147    if let Some(identifier) = error.identifier {
148        builder = builder.with_identifier(identifier);
149    }
150    builder.build()
151}
152
153fn gradient_descriptor_error_with_detail(
154    error: &'static BuiltinErrorDescriptor,
155    detail: impl AsRef<str>,
156) -> RuntimeError {
157    gradient_descriptor_error_with_message(format!("{}: {}", error.message, detail.as_ref()), error)
158}
159
160fn gradient_invalid_argument(detail: impl AsRef<str>) -> RuntimeError {
161    gradient_descriptor_error_with_detail(&GRADIENT_ERROR_INVALID_ARGUMENT, detail)
162}
163
164fn gradient_invalid_input(detail: impl AsRef<str>) -> RuntimeError {
165    gradient_descriptor_error_with_detail(&GRADIENT_ERROR_INVALID_INPUT, detail)
166}
167
168fn gradient_internal_error(detail: impl AsRef<str>) -> RuntimeError {
169    gradient_descriptor_error_with_detail(&GRADIENT_ERROR_INTERNAL, detail)
170}
171
172#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::math::reduction::gradient")]
173pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
174    name: "gradient",
175    op_kind: GpuOpKind::Custom("numerical-gradient"),
176    supported_precisions: &[ScalarType::F32, ScalarType::F64],
177    broadcast: BroadcastSemantics::Matlab,
178    provider_hooks: &[ProviderHook::Custom("gradient_dim")],
179    constant_strategy: ConstantStrategy::InlineLiteral,
180    residency: ResidencyPolicy::NewHandle,
181    nan_mode: ReductionNaN::Include,
182    two_pass_threshold: None,
183    workgroup_size: None,
184    accepts_nan_mode: false,
185    notes:
186        "Providers may keep scalar-spacing gradients on device via `gradient_dim`; coordinate-vector spacing falls back to the host in this implementation.",
187};
188
189#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::math::reduction::gradient")]
190pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
191    name: "gradient",
192    shape: ShapeRequirements::Any,
193    constant_strategy: ConstantStrategy::InlineLiteral,
194    elementwise: None,
195    reduction: None,
196    emits_nan: false,
197    notes: "Gradient preserves input shape and uses edge-aware finite differences, so providers expose it through a custom sink hook.",
198};
199
200#[runtime_builtin(
201    name = "gradient",
202    category = "math/reduction",
203    summary = "Compute numerical gradients.",
204    keywords = "gradient,numerical gradient,finite difference,vector field,gpu",
205    accel = "gradient",
206    type_resolver(gradient_type),
207    descriptor(crate::builtins::math::reduction::gradient::GRADIENT_DESCRIPTOR),
208    builtin_path = "crate::builtins::math::reduction::gradient"
209)]
210async fn gradient_builtin(value: Value, rest: Vec<Value>) -> crate::BuiltinResult<Value> {
211    let requested_outputs = crate::output_count::current_output_count().unwrap_or(1);
212    if requested_outputs == 0 {
213        return Ok(Value::OutputList(Vec::new()));
214    }
215
216    let available_outputs = gradient_output_dims(value_shape(&value), value_len(&value));
217    if requested_outputs > available_outputs.len() {
218        return Err(gradient_invalid_argument(format!(
219            "gradient: requested {requested_outputs} outputs, but input supports at most {}",
220            available_outputs.len()
221        )));
222    }
223
224    let spacings = parse_spacings(&rest, available_outputs.len()).await?;
225    let outputs =
226        evaluate_gradient_outputs(value, &available_outputs[..requested_outputs], &spacings)
227            .await?;
228
229    if crate::output_count::current_output_count().is_some() {
230        return Ok(Value::OutputList(outputs));
231    }
232
233    Ok(outputs
234        .into_iter()
235        .next()
236        .expect("single-output gradient result"))
237}
238
239async fn evaluate_gradient_outputs(
240    value: Value,
241    requested_dims: &[usize],
242    all_spacings: &[f64],
243) -> BuiltinResult<Vec<Value>> {
244    if let Value::GpuTensor(handle) = value {
245        return gradient_gpu_outputs(handle, requested_dims, all_spacings).await;
246    }
247
248    evaluate_host_gradient_outputs(value, requested_dims, all_spacings)
249}
250
251fn evaluate_host_gradient_outputs(
252    value: Value,
253    requested_dims: &[usize],
254    all_spacings: &[f64],
255) -> BuiltinResult<Vec<Value>> {
256    match value {
257        Value::Tensor(tensor) => {
258            let mut outputs = Vec::with_capacity(requested_dims.len());
259            for &dim in requested_dims {
260                let spacing = spacing_for_dim(dim, requested_dims, all_spacings);
261                outputs.push(tensor::tensor_into_value(gradient_real_tensor_host(
262                    tensor.clone(),
263                    dim,
264                    spacing,
265                )?));
266            }
267            Ok(outputs)
268        }
269        Value::LogicalArray(logical) => {
270            let tensor = tensor::logical_to_tensor(&logical).map_err(gradient_invalid_input)?;
271            let mut outputs = Vec::with_capacity(requested_dims.len());
272            for &dim in requested_dims {
273                let spacing = spacing_for_dim(dim, requested_dims, all_spacings);
274                outputs.push(tensor::tensor_into_value(gradient_real_tensor_host(
275                    tensor.clone(),
276                    dim,
277                    spacing,
278                )?));
279            }
280            Ok(outputs)
281        }
282        Value::Num(_) | Value::Int(_) | Value::Bool(_) => {
283            let tensor =
284                tensor::value_into_tensor_for(NAME, value).map_err(gradient_invalid_input)?;
285            let mut outputs = Vec::with_capacity(requested_dims.len());
286            for &dim in requested_dims {
287                let spacing = spacing_for_dim(dim, requested_dims, all_spacings);
288                outputs.push(tensor::tensor_into_value(gradient_real_tensor_host(
289                    tensor.clone(),
290                    dim,
291                    spacing,
292                )?));
293            }
294            Ok(outputs)
295        }
296        Value::Complex(re, im) => {
297            let tensor = ComplexTensor {
298                data: vec![(re, im)],
299                shape: vec![1, 1],
300                rows: 1,
301                cols: 1,
302            };
303            let mut outputs = Vec::with_capacity(requested_dims.len());
304            for &dim in requested_dims {
305                let spacing = spacing_for_dim(dim, requested_dims, all_spacings);
306                outputs.push(complex_tensor_into_value(gradient_complex_tensor_host(
307                    tensor.clone(),
308                    dim,
309                    spacing,
310                )?));
311            }
312            Ok(outputs)
313        }
314        Value::ComplexTensor(tensor) => {
315            let mut outputs = Vec::with_capacity(requested_dims.len());
316            for &dim in requested_dims {
317                let spacing = spacing_for_dim(dim, requested_dims, all_spacings);
318                outputs.push(complex_tensor_into_value(gradient_complex_tensor_host(
319                    tensor.clone(),
320                    dim,
321                    spacing,
322                )?));
323            }
324            Ok(outputs)
325        }
326        other => Err(gradient_invalid_input(format!(
327            "gradient: unsupported input type {:?}; expected numeric or logical data",
328            other
329        ))),
330    }
331}
332
333async fn gradient_gpu_outputs(
334    handle: GpuTensorHandle,
335    requested_dims: &[usize],
336    all_spacings: &[f64],
337) -> BuiltinResult<Vec<Value>> {
338    if runmat_accelerate_api::handle_storage(&handle) == GpuTensorStorage::ComplexInterleaved {
339        let gathered = gpu_helpers::gather_value_async(&Value::GpuTensor(handle)).await?;
340        return evaluate_host_gradient_outputs(gathered, requested_dims, all_spacings);
341    }
342
343    if let Some(provider) = runmat_accelerate_api::provider() {
344        let mut outputs = Vec::with_capacity(requested_dims.len());
345        for &dim in requested_dims {
346            let spacing = spacing_for_dim(dim, requested_dims, all_spacings);
347            match provider.gradient_dim(&handle, dim.saturating_sub(1), spacing) {
348                Ok(device_result) => outputs.push(gpu_helpers::resident_gpu_value(device_result)),
349                Err(_) => {
350                    let gathered =
351                        gpu_helpers::gather_value_async(&Value::GpuTensor(handle)).await?;
352                    return evaluate_host_gradient_outputs(gathered, requested_dims, all_spacings);
353                }
354            }
355        }
356        return Ok(outputs);
357    }
358
359    let gathered = gpu_helpers::gather_value_async(&Value::GpuTensor(handle)).await?;
360    evaluate_host_gradient_outputs(gathered, requested_dims, all_spacings)
361}
362
363fn spacing_for_dim(dim: usize, available_dims: &[usize], spacings: &[f64]) -> f64 {
364    if spacings.len() == 1 {
365        return spacings[0];
366    }
367
368    let index = available_dims
369        .iter()
370        .position(|candidate| *candidate == dim)
371        .expect("spacing lookup requires matching dimension");
372    spacings[index]
373}
374
375async fn parse_spacings(args: &[Value], available_dims: usize) -> BuiltinResult<Vec<f64>> {
376    match args.len() {
377        0 => Ok(vec![1.0; available_dims]),
378        1 => {
379            let spacing = parse_scalar_spacing(&args[0]).await?;
380            Ok(vec![spacing; available_dims])
381        }
382        count if count == available_dims => {
383            let mut spacings = Vec::with_capacity(args.len());
384            for value in args {
385                spacings.push(parse_scalar_spacing(value).await?);
386            }
387            Ok(spacings)
388        }
389        _ => Err(gradient_invalid_argument(format!(
390            "gradient: expected 0, 1, or {available_dims} scalar spacing arguments"
391        ))),
392    }
393}
394
395async fn parse_scalar_spacing(value: &Value) -> BuiltinResult<f64> {
396    match value {
397        Value::Tensor(tensor) if tensor.data.is_empty() => {
398            return Err(gradient_invalid_argument(
399                "gradient: empty spacing arguments are not supported",
400            ))
401        }
402        _ => {}
403    }
404
405    let Some(spacing) = tensor::scalar_f64_from_value_async(value)
406        .await
407        .map_err(gradient_invalid_argument)?
408    else {
409        return Err(gradient_invalid_argument(
410            "gradient: only scalar spacings are supported in this implementation",
411        ));
412    };
413
414    if !spacing.is_finite() {
415        return Err(gradient_invalid_argument(
416            "gradient: spacing must be finite",
417        ));
418    }
419    if spacing == 0.0 {
420        return Err(gradient_invalid_argument(
421            "gradient: spacing must be nonzero",
422        ));
423    }
424    Ok(spacing)
425}
426
427fn value_shape(value: &Value) -> &[usize] {
428    match value {
429        Value::Tensor(tensor) => &tensor.shape,
430        Value::LogicalArray(logical) => &logical.shape,
431        Value::ComplexTensor(tensor) => &tensor.shape,
432        Value::GpuTensor(handle) => &handle.shape,
433        _ => &[],
434    }
435}
436
437fn value_len(value: &Value) -> usize {
438    match value {
439        Value::Tensor(tensor) => tensor.data.len(),
440        Value::LogicalArray(logical) => logical.data.len(),
441        Value::ComplexTensor(tensor) => tensor.data.len(),
442        Value::GpuTensor(handle) => product(&handle.shape),
443        _ => 1,
444    }
445}
446
447pub fn matlab_gradient_shape(shape: &[usize], len: usize) -> Vec<usize> {
448    if shape.is_empty() {
449        if len == 0 {
450            Vec::new()
451        } else {
452            vec![1, 1]
453        }
454    } else if shape.len() == 1 {
455        if shape[0] == 1 {
456            vec![1, 1]
457        } else {
458            vec![1, shape[0]]
459        }
460    } else {
461        shape.to_vec()
462    }
463}
464
465fn gradient_output_dims(shape: &[usize], len: usize) -> Vec<usize> {
466    let normalized_shape = matlab_gradient_shape(shape, len);
467    let mut ext_shape = if normalized_shape.is_empty() {
468        if len == 0 {
469            vec![0, 0]
470        } else {
471            vec![1, 1]
472        }
473    } else {
474        normalized_shape
475    };
476    if ext_shape.len() == 1 {
477        ext_shape.push(1);
478    }
479
480    if ext_shape.len() <= 2 {
481        let rows = ext_shape.first().copied().unwrap_or(1);
482        let cols = ext_shape.get(1).copied().unwrap_or(1);
483        if rows == 1 && cols == 1 {
484            vec![1]
485        } else if rows == 1 {
486            vec![2]
487        } else if cols == 1 {
488            vec![1]
489        } else {
490            vec![2, 1]
491        }
492    } else {
493        let mut dims = vec![2, 1];
494        for dim in 3..=ext_shape.len() {
495            dims.push(dim);
496        }
497        dims
498    }
499}
500
501pub fn gradient_real_tensor_host(
502    tensor: Tensor,
503    dim: usize,
504    spacing: f64,
505) -> BuiltinResult<Tensor> {
506    let Tensor {
507        data, shape, dtype, ..
508    } = tensor;
509    let dim_index = dim.saturating_sub(1);
510    let mut shape = matlab_gradient_shape(&shape, data.len());
511
512    if data.is_empty() {
513        // Return early before the `push(1)` padding loop: that loop would give a
514        // shape like [1] or [1,1] whose product is 1 ≠ 0, violating Tensor's
515        // invariant. Use the normalised shape directly, falling back to [0,0] if
516        // matlab_gradient_shape returned an empty vec (untyped empty tensor).
517        let empty_shape = if shape.is_empty() { vec![0, 0] } else { shape };
518        return Tensor::new_with_dtype(Vec::new(), empty_shape, dtype)
519            .map_err(|e| gradient_internal_error(format!("gradient: {e}")));
520    }
521
522    while shape.len() <= dim_index {
523        shape.push(1);
524    }
525
526    let mut ext_shape = shape.clone();
527    while ext_shape.len() <= dim_index {
528        ext_shape.push(1);
529    }
530    let len_dim = ext_shape[dim_index];
531    let stride_before = if dim_index == 0 {
532        1usize
533    } else {
534        product(&ext_shape[..dim_index]).max(1)
535    };
536    let stride_after = if dim_index + 1 >= ext_shape.len() {
537        1usize
538    } else {
539        product(&ext_shape[dim_index + 1..]).max(1)
540    };
541
542    let mut out = vec![0.0; data.len()];
543    if len_dim > 1 {
544        let block = stride_before
545            .checked_mul(len_dim)
546            .ok_or_else(|| gradient_internal_error("gradient: block size overflow"))?;
547        for after in 0..stride_after {
548            let base = after
549                .checked_mul(block)
550                .ok_or_else(|| gradient_internal_error("gradient: indexing overflow"))?;
551            for before in 0..stride_before {
552                for k in 0..len_dim {
553                    let idx = base + before + k * stride_before;
554                    out[idx] = if k == 0 {
555                        (data[idx + stride_before] - data[idx]) / spacing
556                    } else if k + 1 == len_dim {
557                        (data[idx] - data[idx - stride_before]) / spacing
558                    } else {
559                        (data[idx + stride_before] - data[idx - stride_before]) / (2.0 * spacing)
560                    };
561                }
562            }
563        }
564    }
565
566    Tensor::new_with_dtype(out, shape, dtype)
567        .map_err(|e| gradient_internal_error(format!("gradient: {e}")))
568}
569
570pub fn gradient_complex_tensor_host(
571    tensor: ComplexTensor,
572    dim: usize,
573    spacing: f64,
574) -> BuiltinResult<ComplexTensor> {
575    let ComplexTensor { data, shape, .. } = tensor;
576    let dim_index = dim.saturating_sub(1);
577    let mut shape = matlab_gradient_shape(&shape, data.len());
578
579    if data.is_empty() {
580        // Same fix as gradient_real_tensor_host: avoid padding the shape with 1s
581        // before the early return, which would produce product ≠ 0 for empty data.
582        let empty_shape = if shape.is_empty() { vec![0, 0] } else { shape };
583        return ComplexTensor::new(Vec::new(), empty_shape)
584            .map_err(|e| gradient_internal_error(format!("gradient: {e}")));
585    }
586
587    while shape.len() <= dim_index {
588        shape.push(1);
589    }
590
591    let mut ext_shape = shape.clone();
592    while ext_shape.len() <= dim_index {
593        ext_shape.push(1);
594    }
595    let len_dim = ext_shape[dim_index];
596    let stride_before = if dim_index == 0 {
597        1usize
598    } else {
599        product(&ext_shape[..dim_index]).max(1)
600    };
601    let stride_after = if dim_index + 1 >= ext_shape.len() {
602        1usize
603    } else {
604        product(&ext_shape[dim_index + 1..]).max(1)
605    };
606
607    let mut out = vec![(0.0, 0.0); data.len()];
608    if len_dim > 1 {
609        let block = stride_before
610            .checked_mul(len_dim)
611            .ok_or_else(|| gradient_internal_error("gradient: block size overflow"))?;
612        for after in 0..stride_after {
613            let base = after
614                .checked_mul(block)
615                .ok_or_else(|| gradient_internal_error("gradient: indexing overflow"))?;
616            for before in 0..stride_before {
617                for k in 0..len_dim {
618                    let idx = base + before + k * stride_before;
619                    out[idx] = if k == 0 {
620                        scale_complex(
621                            sub_complex(data[idx + stride_before], data[idx]),
622                            1.0 / spacing,
623                        )
624                    } else if k + 1 == len_dim {
625                        scale_complex(
626                            sub_complex(data[idx], data[idx - stride_before]),
627                            1.0 / spacing,
628                        )
629                    } else {
630                        scale_complex(
631                            sub_complex(data[idx + stride_before], data[idx - stride_before]),
632                            0.5 / spacing,
633                        )
634                    };
635                }
636            }
637        }
638    }
639
640    ComplexTensor::new(out, shape).map_err(|e| gradient_internal_error(format!("gradient: {e}")))
641}
642
643fn sub_complex(lhs: (f64, f64), rhs: (f64, f64)) -> (f64, f64) {
644    (lhs.0 - rhs.0, lhs.1 - rhs.1)
645}
646
647fn scale_complex(value: (f64, f64), scale: f64) -> (f64, f64) {
648    (value.0 * scale, value.1 * scale)
649}
650
651fn product(dims: &[usize]) -> usize {
652    dims.iter()
653        .copied()
654        .fold(1usize, |acc, value| acc.saturating_mul(value))
655}
656
657#[cfg(test)]
658mod tests {
659    use super::*;
660    use crate::builtins::common::test_support;
661    use futures::executor::block_on;
662    #[cfg(feature = "wgpu")]
663    use runmat_accelerate_api::AccelProvider;
664    #[cfg(feature = "wgpu")]
665    use runmat_accelerate_api::HostTensorView;
666    use runmat_builtins::{NumericDType, Tensor};
667
668    fn gradient_builtin(value: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
669        block_on(super::gradient_builtin(value, rest))
670    }
671
672    #[test]
673    fn gradient_descriptor_signatures_cover_core_forms() {
674        let labels: Vec<&str> = GRADIENT_DESCRIPTOR
675            .signatures
676            .iter()
677            .map(|sig| sig.label)
678            .collect();
679        assert!(labels.contains(&"G = gradient(F)"));
680        assert!(labels.contains(&"G = gradient(F, h)"));
681        assert!(labels.contains(&"[G1, G2, ...] = gradient(F)"));
682        assert!(labels.contains(&"[G1, G2, ...] = gradient(F, h1, h2, ...)"));
683    }
684
685    #[test]
686    fn gradient_descriptor_errors_have_stable_codes() {
687        assert!(GRADIENT_DESCRIPTOR
688            .errors
689            .iter()
690            .any(|error| error.code == GRADIENT_ERROR_INVALID_ARGUMENT.code));
691        assert!(GRADIENT_DESCRIPTOR
692            .errors
693            .iter()
694            .any(|error| error.code == GRADIENT_ERROR_INVALID_INPUT.code));
695        assert!(GRADIENT_DESCRIPTOR
696            .errors
697            .iter()
698            .any(|error| error.code == GRADIENT_ERROR_INTERNAL.code));
699    }
700
701    #[test]
702    fn gradient_row_vector_returns_horizontal_derivative() {
703        let tensor = Tensor::new(vec![1.0, 4.0, 9.0], vec![1, 3]).unwrap();
704        let result = gradient_builtin(Value::Tensor(tensor), Vec::new()).expect("gradient");
705        assert_eq!(
706            result,
707            Value::Tensor(Tensor::new(vec![3.0, 4.0, 5.0], vec![1, 3]).unwrap())
708        );
709    }
710
711    #[test]
712    fn gradient_one_dimensional_tensor_is_treated_as_row_vector() {
713        let tensor = Tensor::new(vec![1.0, 4.0, 9.0], vec![3]).unwrap();
714        let result =
715            gradient_builtin(Value::Tensor(tensor), vec![Value::Num(2.0)]).expect("gradient");
716        match result {
717            Value::Tensor(out) => {
718                assert_eq!(out.shape, vec![1, 3]);
719                assert_eq!(out.data, vec![1.5, 2.0, 2.5]);
720            }
721            other => panic!("expected tensor, got {other:?}"),
722        }
723    }
724
725    #[test]
726    fn gradient_matrix_outputs_follow_matlab_order() {
727        let tensor = Tensor::new(vec![1.0, 3.0, 2.0, 4.0], vec![2, 2]).unwrap();
728        let _guard = crate::output_count::push_output_count(Some(2));
729        let result = gradient_builtin(Value::Tensor(tensor), Vec::new()).expect("gradient");
730        match result {
731            Value::OutputList(outputs) => {
732                let fx = test_support::gather(outputs[0].clone()).expect("fx");
733                let fy = test_support::gather(outputs[1].clone()).expect("fy");
734                assert_eq!(fx.data, vec![1.0, 1.0, 1.0, 1.0]);
735                assert_eq!(fy.data, vec![2.0, 2.0, 2.0, 2.0]);
736            }
737            other => panic!("expected output list, got {other:?}"),
738        }
739    }
740
741    #[test]
742    fn gradient_scalar_spacing_scales_output() {
743        let tensor = Tensor::new(vec![1.0, 4.0, 9.0], vec![1, 3]).unwrap();
744        let result =
745            gradient_builtin(Value::Tensor(tensor), vec![Value::Num(2.0)]).expect("gradient");
746        match result {
747            Value::Tensor(out) => assert_eq!(out.data, vec![1.5, 2.0, 2.5]),
748            other => panic!("expected tensor, got {other:?}"),
749        }
750    }
751
752    #[test]
753    fn gradient_preserves_single_precision_host_tensor() {
754        let tensor =
755            Tensor::new_with_dtype(vec![1.0, 4.0, 9.0], vec![1, 3], NumericDType::F32).unwrap();
756        let result = gradient_builtin(Value::Tensor(tensor), Vec::new()).expect("gradient");
757        match result {
758            Value::Tensor(out) => assert_eq!(out.dtype, NumericDType::F32),
759            other => panic!("expected tensor, got {other:?}"),
760        }
761    }
762
763    #[test]
764    fn gradient_complex_host_supported() {
765        let tensor =
766            ComplexTensor::new(vec![(1.0, 1.0), (4.0, 3.0), (9.0, 6.0)], vec![1, 3]).unwrap();
767        let result = gradient_builtin(Value::ComplexTensor(tensor), Vec::new()).expect("gradient");
768        match result {
769            Value::ComplexTensor(out) => {
770                assert_eq!(out.data, vec![(3.0, 2.0), (4.0, 2.5), (5.0, 3.0)]);
771            }
772            other => panic!("expected complex tensor, got {other:?}"),
773        }
774    }
775
776    #[test]
777    fn gradient_rejects_coordinate_vector_spacing_in_v1() {
778        let tensor = Tensor::new(vec![1.0, 4.0, 9.0], vec![1, 3]).unwrap();
779        let spacing = Tensor::new(vec![0.0, 1.0, 2.0], vec![1, 3]).unwrap();
780        let err =
781            gradient_builtin(Value::Tensor(tensor), vec![Value::Tensor(spacing)]).unwrap_err();
782        assert_eq!(err.identifier(), GRADIENT_ERROR_INVALID_ARGUMENT.identifier);
783        assert!(err.message().contains("scalar"));
784    }
785
786    #[test]
787    fn gradient_rejects_too_many_outputs() {
788        let tensor = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
789        let _guard = crate::output_count::push_output_count(Some(2));
790        let err = gradient_builtin(Value::Tensor(tensor), Vec::new()).unwrap_err();
791        assert_eq!(err.identifier(), GRADIENT_ERROR_INVALID_ARGUMENT.identifier);
792        assert!(err.message().contains("requested 2 outputs"));
793    }
794
795    #[test]
796    #[cfg(feature = "wgpu")]
797    fn gradient_gpu_scalar_spacing_matches_cpu_and_stays_resident() {
798        let _guard = test_support::accel_test_lock();
799        let Ok(provider) = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
800            runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
801        ) else {
802            return;
803        };
804        let host =
805            Tensor::new_with_dtype(vec![1.0, 4.0, 9.0], vec![1, 3], NumericDType::F32).unwrap();
806        let view = HostTensorView {
807            data: &host.data,
808            shape: &host.shape,
809        };
810        let handle = provider.upload(&view).expect("upload");
811        let result =
812            gradient_builtin(Value::GpuTensor(handle), vec![Value::Num(2.0)]).expect("gradient");
813        match result {
814            Value::GpuTensor(out) => {
815                let gathered = test_support::gather(Value::GpuTensor(out)).expect("gather");
816                assert_eq!(gathered.data, vec![1.5, 2.0, 2.5]);
817                assert_eq!(gathered.dtype, NumericDType::F32);
818            }
819            other => panic!("expected gpu tensor, got {other:?}"),
820        }
821    }
822
823    #[test]
824    #[cfg(feature = "wgpu")]
825    fn gradient_gpu_one_dimensional_shape_matches_matlab_row_vector_semantics() {
826        let _guard = test_support::accel_test_lock();
827        let Ok(provider) = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
828            runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
829        ) else {
830            return;
831        };
832        let data = [1.0, 4.0, 9.0];
833        let shape = [3usize];
834        let view = HostTensorView {
835            data: &data,
836            shape: &shape,
837        };
838        let handle = provider.upload(&view).expect("upload");
839        let result =
840            gradient_builtin(Value::GpuTensor(handle), vec![Value::Num(2.0)]).expect("gradient");
841        let gathered = test_support::gather(result).expect("gather");
842        assert_eq!(gathered.shape, vec![1, 3]);
843        assert_eq!(gathered.data, vec![1.5, 2.0, 2.5]);
844    }
845
846    #[test]
847    #[cfg(feature = "wgpu")]
848    fn gradient_gpu_multi_output_uses_output_list() {
849        let _guard = test_support::accel_test_lock();
850        let Ok(provider) = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
851            runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
852        ) else {
853            return;
854        };
855        let host = Tensor::new(vec![1.0, 3.0, 2.0, 4.0], vec![2, 2]).unwrap();
856        let view = HostTensorView {
857            data: &host.data,
858            shape: &host.shape,
859        };
860        let handle = provider.upload(&view).expect("upload");
861        let _out_guard = crate::output_count::push_output_count(Some(2));
862        let result = gradient_builtin(Value::GpuTensor(handle), Vec::new()).expect("gradient");
863        match result {
864            Value::OutputList(outputs) => {
865                assert!(matches!(outputs[0], Value::GpuTensor(_)));
866                assert!(matches!(outputs[1], Value::GpuTensor(_)));
867            }
868            other => panic!("expected output list, got {other:?}"),
869        }
870    }
871}