1use runmat_accelerate_api::{
4 AccelProvider, CovNormalization, CovRows, CovarianceOptions, GpuTensorHandle, GpuTensorStorage,
5 HostTensorView,
6};
7use runmat_builtins::{
8 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
9 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
10 Tensor, Value,
11};
12use runmat_macros::runtime_builtin;
13
14use crate::builtins::common::gpu_helpers;
15use crate::builtins::common::spec::{
16 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
17 ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
18};
19use crate::builtins::common::tensor;
20use crate::builtins::stats::type_resolvers::cov_type;
21use crate::{build_runtime_error, BuiltinResult, RuntimeError};
22
23const NAME: &str = "cov";
24const COV_OUTPUT: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
25 name: "C",
26 ty: BuiltinParamType::NumericArray,
27 arity: BuiltinParamArity::Required,
28 default: None,
29 description: "Covariance matrix.",
30}];
31
32const COV_INPUTS_X: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
33 name: "X",
34 ty: BuiltinParamType::Any,
35 arity: BuiltinParamArity::Required,
36 default: None,
37 description: "Input observations (rows are observations, columns are variables).",
38}];
39
40const COV_INPUTS_X_Y_OR_W: [BuiltinParamDescriptor; 2] = [
41 BuiltinParamDescriptor {
42 name: "X",
43 ty: BuiltinParamType::Any,
44 arity: BuiltinParamArity::Required,
45 default: None,
46 description: "Input observations (rows are observations, columns are variables).",
47 },
48 BuiltinParamDescriptor {
49 name: "Y_or_w",
50 ty: BuiltinParamType::Any,
51 arity: BuiltinParamArity::Required,
52 default: None,
53 description: "Second dataset (Y) or weight vector (w), depending on shape/position.",
54 },
55];
56
57const COV_INPUTS_X_NORMALIZATION: [BuiltinParamDescriptor; 2] = [
58 BuiltinParamDescriptor {
59 name: "X",
60 ty: BuiltinParamType::Any,
61 arity: BuiltinParamArity::Required,
62 default: None,
63 description: "Input observations (rows are observations, columns are variables).",
64 },
65 BuiltinParamDescriptor {
66 name: "normalization",
67 ty: BuiltinParamType::NumericScalar,
68 arity: BuiltinParamArity::Required,
69 default: Some("0"),
70 description: "Normalization flag: 0 (unbiased) or 1 (biased).",
71 },
72];
73
74const COV_INPUTS_X_ROWS: [BuiltinParamDescriptor; 2] = [
75 BuiltinParamDescriptor {
76 name: "X",
77 ty: BuiltinParamType::Any,
78 arity: BuiltinParamArity::Required,
79 default: None,
80 description: "Input observations (rows are observations, columns are variables).",
81 },
82 BuiltinParamDescriptor {
83 name: "rows_option",
84 ty: BuiltinParamType::StringScalar,
85 arity: BuiltinParamArity::Required,
86 default: Some("\"all\""),
87 description: "Rows handling mode: 'all', 'omitrows', or 'partialrows'.",
88 },
89];
90
91const COV_INPUTS_X_Y_OPT: [BuiltinParamDescriptor; 3] = [
92 BuiltinParamDescriptor {
93 name: "X",
94 ty: BuiltinParamType::Any,
95 arity: BuiltinParamArity::Required,
96 default: None,
97 description: "Input observations (rows are observations, columns are variables).",
98 },
99 BuiltinParamDescriptor {
100 name: "Y",
101 ty: BuiltinParamType::Any,
102 arity: BuiltinParamArity::Required,
103 default: None,
104 description: "Second dataset with matching row count.",
105 },
106 BuiltinParamDescriptor {
107 name: "opt",
108 ty: BuiltinParamType::Any,
109 arity: BuiltinParamArity::Required,
110 default: None,
111 description: "Normalization flag or rows option.",
112 },
113];
114
115const COV_INPUTS_X_Y_W: [BuiltinParamDescriptor; 3] = [
116 BuiltinParamDescriptor {
117 name: "X",
118 ty: BuiltinParamType::Any,
119 arity: BuiltinParamArity::Required,
120 default: None,
121 description: "Input observations (rows are observations, columns are variables).",
122 },
123 BuiltinParamDescriptor {
124 name: "Y",
125 ty: BuiltinParamType::Any,
126 arity: BuiltinParamArity::Required,
127 default: None,
128 description: "Second dataset with matching row count.",
129 },
130 BuiltinParamDescriptor {
131 name: "w",
132 ty: BuiltinParamType::Any,
133 arity: BuiltinParamArity::Required,
134 default: None,
135 description: "Weight vector with one weight per observation row.",
136 },
137];
138
139const COV_INPUTS_X_Y_W_OPT: [BuiltinParamDescriptor; 4] = [
140 BuiltinParamDescriptor {
141 name: "X",
142 ty: BuiltinParamType::Any,
143 arity: BuiltinParamArity::Required,
144 default: None,
145 description: "Input observations (rows are observations, columns are variables).",
146 },
147 BuiltinParamDescriptor {
148 name: "Y",
149 ty: BuiltinParamType::Any,
150 arity: BuiltinParamArity::Required,
151 default: None,
152 description: "Second dataset with matching row count.",
153 },
154 BuiltinParamDescriptor {
155 name: "w",
156 ty: BuiltinParamType::Any,
157 arity: BuiltinParamArity::Required,
158 default: None,
159 description: "Weight vector with one weight per observation row.",
160 },
161 BuiltinParamDescriptor {
162 name: "opt",
163 ty: BuiltinParamType::Any,
164 arity: BuiltinParamArity::Required,
165 default: None,
166 description: "Normalization flag or rows option.",
167 },
168];
169
170const COV_SIGNATURES: [BuiltinSignatureDescriptor; 7] = [
171 BuiltinSignatureDescriptor {
172 label: "C = cov(X)",
173 inputs: &COV_INPUTS_X,
174 outputs: &COV_OUTPUT,
175 },
176 BuiltinSignatureDescriptor {
177 label: "C = cov(X, Y_or_w)",
178 inputs: &COV_INPUTS_X_Y_OR_W,
179 outputs: &COV_OUTPUT,
180 },
181 BuiltinSignatureDescriptor {
182 label: "C = cov(X, normalization)",
183 inputs: &COV_INPUTS_X_NORMALIZATION,
184 outputs: &COV_OUTPUT,
185 },
186 BuiltinSignatureDescriptor {
187 label: "C = cov(X, rows_option)",
188 inputs: &COV_INPUTS_X_ROWS,
189 outputs: &COV_OUTPUT,
190 },
191 BuiltinSignatureDescriptor {
192 label: "C = cov(X, Y, opt)",
193 inputs: &COV_INPUTS_X_Y_OPT,
194 outputs: &COV_OUTPUT,
195 },
196 BuiltinSignatureDescriptor {
197 label: "C = cov(X, Y, w)",
198 inputs: &COV_INPUTS_X_Y_W,
199 outputs: &COV_OUTPUT,
200 },
201 BuiltinSignatureDescriptor {
202 label: "C = cov(X, Y, w, opt)",
203 inputs: &COV_INPUTS_X_Y_W_OPT,
204 outputs: &COV_OUTPUT,
205 },
206];
207
208const COV_ERROR_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
209 code: "RM.COV.INVALID_ARGUMENT",
210 identifier: Some("RunMat:cov:InvalidArgument"),
211 when: "Arguments are malformed or unsupported for cov.",
212 message: "cov: invalid argument",
213};
214
215const COV_ERROR_COMPLEX_UNSUPPORTED: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
216 code: "RM.COV.COMPLEX_UNSUPPORTED",
217 identifier: Some("RunMat:cov:ComplexUnsupported"),
218 when: "Any argument is complex-valued.",
219 message: "cov: complex inputs are not supported yet",
220};
221
222const COV_ERROR_ROWS_MISMATCH: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
223 code: "RM.COV.ROWS_MISMATCH",
224 identifier: Some("RunMat:cov:RowsMismatch"),
225 when: "Two input datasets do not have the same number of rows.",
226 message: "cov: inputs must have the same number of rows",
227};
228
229const COV_ERROR_NORMALIZATION_INVALID: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
230 code: "RM.COV.NORMALIZATION_INVALID",
231 identifier: Some("RunMat:cov:NormalizationInvalid"),
232 when: "Normalization flag is non-finite, non-integer, or not 0/1.",
233 message: "cov: normalization flag is invalid",
234};
235
236const COV_ERROR_WEIGHT_VECTOR_LENGTH_MISMATCH: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
237 code: "RM.COV.WEIGHT_VECTOR_LENGTH_MISMATCH",
238 identifier: Some("RunMat:cov:WeightVectorLengthMismatch"),
239 when: "Weight vector length does not match observation row count.",
240 message: "cov: weight vector length mismatch",
241};
242
243const COV_ERROR_ROWS_OPTION_UNKNOWN: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
244 code: "RM.COV.ROWS_OPTION_UNKNOWN",
245 identifier: Some("RunMat:cov:RowsOptionUnknown"),
246 when: "Rows option is not one of all/omitrows/partialrows.",
247 message: "cov: unknown rows option",
248};
249
250const COV_ERROR_NORMALIZATION_DUPLICATE: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
251 code: "RM.COV.NORMALIZATION_DUPLICATE",
252 identifier: Some("RunMat:cov:NormalizationDuplicate"),
253 when: "Normalization flag is provided more than once.",
254 message: "cov: normalization flag specified more than once",
255};
256
257const COV_ERROR_TOO_MANY_ARRAY_ARGUMENTS: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
258 code: "RM.COV.TOO_MANY_ARRAY_ARGUMENTS",
259 identifier: Some("RunMat:cov:TooManyArrayArguments"),
260 when: "More than two data arrays (or Y plus weight) are provided.",
261 message: "cov: too many array arguments",
262};
263
264const COV_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
265 code: "RM.COV.INTERNAL",
266 identifier: Some("RunMat:cov:Internal"),
267 when: "Internal tensor conversion/allocation or covariance computation fails.",
268 message: "cov: internal operation failed",
269};
270
271const COV_ERRORS: [BuiltinErrorDescriptor; 9] = [
272 COV_ERROR_INVALID_ARGUMENT,
273 COV_ERROR_COMPLEX_UNSUPPORTED,
274 COV_ERROR_ROWS_MISMATCH,
275 COV_ERROR_NORMALIZATION_INVALID,
276 COV_ERROR_WEIGHT_VECTOR_LENGTH_MISMATCH,
277 COV_ERROR_ROWS_OPTION_UNKNOWN,
278 COV_ERROR_NORMALIZATION_DUPLICATE,
279 COV_ERROR_TOO_MANY_ARRAY_ARGUMENTS,
280 COV_ERROR_INTERNAL,
281];
282
283pub const COV_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
284 signatures: &COV_SIGNATURES,
285 output_mode: BuiltinOutputMode::Fixed,
286 completion_policy: BuiltinCompletionPolicy::Public,
287 errors: &COV_ERRORS,
288};
289
290fn cov_error_with(
291 error: &'static BuiltinErrorDescriptor,
292 message: impl Into<String>,
293) -> RuntimeError {
294 let mut builder = build_runtime_error(message).with_builtin(NAME);
295 if let Some(identifier) = error.identifier {
296 builder = builder.with_identifier(identifier);
297 }
298 builder.build()
299}
300
301fn cov_error(error: &'static BuiltinErrorDescriptor) -> RuntimeError {
302 cov_error_with(error, error.message)
303}
304
305fn cov_error_with_detail(
306 error: &'static BuiltinErrorDescriptor,
307 detail: impl std::fmt::Display,
308) -> RuntimeError {
309 cov_error_with(error, format!("{}: {detail}", error.message))
310}
311
312fn cov_internal_error(message: impl Into<String>) -> RuntimeError {
313 cov_error_with(&COV_ERROR_INTERNAL, message)
314}
315
316#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::stats::summary::cov")]
317pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
318 name: "cov",
319 op_kind: GpuOpKind::Custom("summary-stats"),
320 supported_precisions: &[ScalarType::F32, ScalarType::F64],
321 broadcast: BroadcastSemantics::None,
322 provider_hooks: &[ProviderHook::Custom("covariance")],
323 constant_strategy: ConstantStrategy::InlineLiteral,
324 residency: ResidencyPolicy::NewHandle,
325 nan_mode: ReductionNaN::Include,
326 two_pass_threshold: None,
327 workgroup_size: None,
328 accepts_nan_mode: false,
329 notes: "GPU execution is available when rows='all' and no weight vector is supplied; other cases fall back to the CPU path.",
330};
331
332#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::stats::summary::cov")]
333pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
334 name: "cov",
335 shape: ShapeRequirements::Any,
336 constant_strategy: ConstantStrategy::InlineLiteral,
337 elementwise: None,
338 reduction: None,
339 emits_nan: true,
340 notes: "The covariance builtin is treated as a fusion boundary and executes via dedicated kernels or the host reference.",
341};
342
343#[runtime_builtin(
344 name = "cov",
345 category = "stats/summary",
346 summary = "Compute covariance matrices.",
347 keywords = "cov,covariance,statistics,weights,gpu",
348 accel = "reduction",
349 type_resolver(cov_type),
350 descriptor(crate::builtins::stats::summary::cov::COV_DESCRIPTOR),
351 builtin_path = "crate::builtins::stats::summary::cov"
352)]
353async fn cov_builtin(value: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
354 let args = CovArgs::parse(value, rest)?;
355 if let Some(result) = cov_try_gpu(&args).await? {
356 return Ok(result);
357 }
358 cov_host(args).await
359}
360
361pub fn cov_from_tensors(
363 left: Tensor,
364 right: Option<Tensor>,
365 rows: CovRows,
366 weight: CovWeightSpec,
367) -> BuiltinResult<Tensor> {
368 let matrix = combine_tensors(left, right)?;
369 if let CovWeightSpec::Vector(ref vec) = weight {
370 if matrix.rows != vec.len() {
371 return Err(cov_error_with_detail(
372 &COV_ERROR_WEIGHT_VECTOR_LENGTH_MISMATCH,
373 format!("expected {} elements", matrix.rows),
374 ));
375 }
376 }
377 match rows {
378 CovRows::All => covariance_dense(&matrix, &weight),
379 CovRows::OmitRows => {
380 let (filtered, filtered_weight) = filter_complete_rows(&matrix, weight);
381 covariance_dense(&filtered, &filtered_weight)
382 }
383 CovRows::PartialRows => covariance_pairwise(&matrix, &weight),
384 }
385}
386
387#[derive(Debug)]
388struct CovArgs {
389 first: Value,
390 second: Option<Value>,
391 normalization: CovNormalization,
392 rows: CovRows,
393 weight_vector: Option<Value>,
394}
395
396impl CovArgs {
397 fn parse(first: Value, rest: Vec<Value>) -> BuiltinResult<Self> {
398 let mut second_candidate: Option<Value> = None;
399 let mut weight_candidate: Option<Value> = None;
400 let mut normalization = CovNormalization::Unbiased;
401 let mut normalization_explicit = false;
402 let mut rows = CovRows::All;
403
404 let iter = rest.into_iter();
405 for arg in iter {
406 match arg {
407 Value::String(_) | Value::StringArray(_) | Value::CharArray(_) => {
408 let key = tensor::value_to_string(&arg)
409 .ok_or_else(|| cov_error(&COV_ERROR_INVALID_ARGUMENT))?;
410 let lowered = key.trim().to_ascii_lowercase();
411 rows = parse_rows_option(&lowered)?;
412 }
413 Value::Tensor(_) | Value::LogicalArray(_) | Value::GpuTensor(_) => {
414 if second_candidate.is_none() {
415 second_candidate = Some(arg);
416 } else if weight_candidate.is_none() {
417 weight_candidate = Some(arg);
418 } else {
419 return Err(cov_error(&COV_ERROR_TOO_MANY_ARRAY_ARGUMENTS));
420 }
421 }
422 Value::Num(_) | Value::Int(_) | Value::Bool(_) => {
423 if normalization_explicit || weight_candidate.is_some() {
424 return Err(cov_error(&COV_ERROR_NORMALIZATION_DUPLICATE));
425 }
426 normalization = parse_normalization(arg)?;
427 normalization_explicit = true;
428 }
429 Value::ComplexTensor(_) => {
430 return Err(cov_error(&COV_ERROR_COMPLEX_UNSUPPORTED));
431 }
432 other => {
433 return Err(cov_error_with_detail(
434 &COV_ERROR_INVALID_ARGUMENT,
435 format!("{other:?}"),
436 ))
437 }
438 }
439 }
440
441 if let Some(weight_array) = weight_candidate {
442 return Ok(Self {
444 first,
445 second: second_candidate,
446 normalization,
447 rows,
448 weight_vector: Some(weight_array),
449 });
450 }
451
452 let mut second = second_candidate;
453 let mut weight_vector: Option<Value> = None;
454
455 if let Some(candidate) = second.take() {
456 if should_treat_as_weight(&first, &candidate, normalization_explicit, rows)? {
457 weight_vector = Some(candidate);
458 } else {
459 second = Some(candidate);
460 }
461 }
462
463 Ok(Self {
464 first,
465 second,
466 normalization,
467 rows,
468 weight_vector,
469 })
470 }
471}
472
473#[derive(Debug, Clone)]
474pub enum CovWeightSpec {
475 Scalar(CovNormalization),
476 Vector(Vec<f64>),
477}
478
479async fn cov_try_gpu(args: &CovArgs) -> BuiltinResult<Option<Value>> {
480 if args.rows != CovRows::All {
481 return Ok(None);
482 }
483
484 let first_handle = match &args.first {
485 Value::GpuTensor(handle) => handle,
486 _ => return Ok(None),
487 };
488
489 let provider = match runmat_accelerate_api::provider_for_handle(first_handle)
490 .or_else(runmat_accelerate_api::provider)
491 {
492 Some(p) => p,
493 None => return Ok(None),
494 };
495
496 let maybe_second_handle = match &args.second {
497 Some(Value::GpuTensor(handle)) => {
498 let Some(second_provider) = runmat_accelerate_api::provider_for_handle(handle) else {
499 return Ok(None);
500 };
501 if !std::ptr::eq(provider, second_provider) {
502 return Ok(None);
503 }
504 Some(handle)
505 }
506 Some(_) => return Ok(None),
507 None => None,
508 };
509
510 let rows = gpu_rows(first_handle)?;
511 let mut temporary_inputs = Vec::new();
512 let weight_handle = match materialize_gpu_weight_vector(
513 provider,
514 args.weight_vector.as_ref(),
515 rows,
516 &mut temporary_inputs,
517 )
518 .await
519 {
520 Ok(weight) => weight,
521 Err(err) => {
522 free_temporary_gpu_inputs(provider, temporary_inputs);
523 return Err(err);
524 }
525 };
526 if args.weight_vector.is_some() && weight_handle.is_none() {
527 free_temporary_gpu_inputs(provider, temporary_inputs);
528 return Ok(None);
529 }
530
531 let options = CovarianceOptions {
532 normalization: args.normalization,
533 rows: args.rows,
534 has_weight_vector: weight_handle.is_some(),
535 };
536
537 match provider
538 .covariance(
539 first_handle,
540 maybe_second_handle,
541 weight_handle.as_ref(),
542 &options,
543 )
544 .await
545 {
546 Ok(result) => {
547 free_temporary_gpu_inputs(provider, temporary_inputs);
548 Ok(Some(Value::GpuTensor(result)))
549 }
550 Err(_) => {
551 free_temporary_gpu_inputs(provider, temporary_inputs);
552 Ok(None)
553 }
554 }
555}
556
557fn gpu_rows(handle: &GpuTensorHandle) -> BuiltinResult<usize> {
558 if handle.shape.len() > 2 {
559 return Err(cov_error_with_detail(
560 &COV_ERROR_INVALID_ARGUMENT,
561 "inputs must be 2-D matrices or vectors",
562 ));
563 }
564 Ok(if handle.shape.is_empty() {
565 1
566 } else {
567 handle.shape[0]
568 })
569}
570
571async fn materialize_gpu_weight_vector(
572 provider: &dyn AccelProvider,
573 value: Option<&Value>,
574 expected_rows: usize,
575 temporary_inputs: &mut Vec<GpuTensorHandle>,
576) -> BuiltinResult<Option<GpuTensorHandle>> {
577 let Some(value) = value else {
578 return Ok(None);
579 };
580 if expected_rows == 0 {
581 return Err(cov_error_with_detail(
582 &COV_ERROR_INVALID_ARGUMENT,
583 "weight vector cannot be empty",
584 ));
585 }
586
587 match value {
588 Value::GpuTensor(handle) => {
589 if runmat_accelerate_api::handle_storage(handle) == GpuTensorStorage::ComplexInterleaved
590 {
591 return Ok(None);
592 }
593 validate_gpu_weight_shape(handle, expected_rows)?;
594 let Some(weight_provider) = runmat_accelerate_api::provider_for_handle(handle) else {
595 return Ok(None);
596 };
597 if !std::ptr::eq(provider, weight_provider) {
598 return Ok(None);
599 }
600 Ok(Some(handle.clone()))
601 }
602 other => {
603 let weights = value_to_weight_vector(other.clone(), expected_rows).await?;
604 let shape = [expected_rows, 1];
605 let handle = provider
606 .upload(&HostTensorView {
607 data: &weights,
608 shape: &shape,
609 })
610 .map_err(|err| cov_internal_error(err.to_string()))?;
611 temporary_inputs.push(handle.clone());
612 Ok(Some(handle))
613 }
614 }
615}
616
617fn validate_gpu_weight_shape(handle: &GpuTensorHandle, expected_rows: usize) -> BuiltinResult<()> {
618 if handle.shape.len() > 2 {
619 return Err(cov_error_with_detail(
620 &COV_ERROR_INVALID_ARGUMENT,
621 "weight vector must be one-dimensional",
622 ));
623 }
624 let rows = if handle.shape.is_empty() {
625 1
626 } else {
627 handle.shape[0]
628 };
629 let cols = if handle.shape.len() >= 2 {
630 handle.shape[1]
631 } else {
632 1
633 };
634 if rows != 1 && cols != 1 {
635 return Err(cov_error_with_detail(
636 &COV_ERROR_INVALID_ARGUMENT,
637 "weight vector must be one-dimensional",
638 ));
639 }
640 if rows != expected_rows && cols != expected_rows {
641 return Err(cov_error_with_detail(
642 &COV_ERROR_WEIGHT_VECTOR_LENGTH_MISMATCH,
643 format!("expected {expected_rows} elements"),
644 ));
645 }
646 Ok(())
647}
648
649fn free_temporary_gpu_inputs(provider: &dyn AccelProvider, handles: Vec<GpuTensorHandle>) {
650 for handle in handles {
651 let _ = provider.free(&handle);
652 }
653}
654
655async fn cov_host(args: CovArgs) -> BuiltinResult<Value> {
656 let CovArgs {
657 first,
658 second,
659 normalization,
660 rows,
661 weight_vector,
662 } = args;
663
664 let left = value_to_tensor_gather(first).await?;
665 let right = match second {
666 Some(value) => Some(value_to_tensor_gather(value).await?),
667 None => None,
668 };
669
670 let weight_spec = if let Some(weight_value) = weight_vector {
671 let vector = value_to_weight_vector(weight_value, left.rows()).await?;
672 CovWeightSpec::Vector(vector)
673 } else {
674 CovWeightSpec::Scalar(normalization)
675 };
676
677 let tensor = cov_from_tensors(left, right, rows, weight_spec)?;
678 Ok(Value::Tensor(tensor))
679}
680
681async fn value_to_tensor_gather(value: Value) -> BuiltinResult<Tensor> {
682 match value {
683 Value::GpuTensor(handle) => gpu_helpers::gather_tensor_async(&handle).await,
684 Value::LogicalArray(logical) => {
685 tensor::logical_to_tensor(&logical).map_err(cov_internal_error)
686 }
687 other => tensor::value_into_tensor_for("cov", other).map_err(cov_internal_error),
688 }
689}
690
691async fn value_to_weight_vector(value: Value, expected_rows: usize) -> BuiltinResult<Vec<f64>> {
692 let tensor = match value {
693 Value::GpuTensor(handle) => gpu_helpers::gather_tensor_async(&handle).await?,
694 Value::LogicalArray(logical) => {
695 tensor::logical_to_tensor(&logical).map_err(cov_internal_error)?
696 }
697 other => tensor::value_into_tensor_for("cov", other).map_err(cov_internal_error)?,
698 };
699
700 if tensor.shape.len() > 2 {
701 return Err(cov_error_with_detail(
702 &COV_ERROR_INVALID_ARGUMENT,
703 "weight vector must be one-dimensional",
704 ));
705 }
706 if tensor.rows() != expected_rows && tensor.cols() != expected_rows {
707 return Err(cov_error_with_detail(
708 &COV_ERROR_WEIGHT_VECTOR_LENGTH_MISMATCH,
709 format!("expected {expected_rows} elements"),
710 ));
711 }
712 for (idx, weight) in tensor.data.iter().enumerate() {
713 if !weight.is_finite() || *weight < 0.0 {
714 return Err(cov_error_with_detail(
715 &COV_ERROR_INVALID_ARGUMENT,
716 format!("weights must be non-negative finite values (index {idx})"),
717 ));
718 }
719 }
720 if tensor.data.is_empty() {
721 return Err(cov_error_with_detail(
722 &COV_ERROR_INVALID_ARGUMENT,
723 "weight vector cannot be empty",
724 ));
725 }
726 Ok(tensor.data)
727}
728
729fn parse_rows_option(value: &str) -> BuiltinResult<CovRows> {
730 match value {
731 "all" => Ok(CovRows::All),
732 "omitrows" | "omit" => Ok(CovRows::OmitRows),
733 "partialrows" | "partial" | "pairwise" => Ok(CovRows::PartialRows),
734 other => Err(cov_error_with_detail(
735 &COV_ERROR_ROWS_OPTION_UNKNOWN,
736 format!("'{other}'"),
737 )),
738 }
739}
740
741fn parse_normalization(value: Value) -> BuiltinResult<CovNormalization> {
742 match value {
743 Value::Int(i) => match i.to_i64() {
744 0 => Ok(CovNormalization::Unbiased),
745 1 => Ok(CovNormalization::Biased),
746 other => Err(cov_error_with_detail(
747 &COV_ERROR_NORMALIZATION_INVALID,
748 format!("expected 0 or 1, received {other}"),
749 )),
750 },
751 Value::Num(n) => {
752 if !n.is_finite() {
753 return Err(cov_error_with_detail(
754 &COV_ERROR_NORMALIZATION_INVALID,
755 "value must be finite",
756 ));
757 }
758 let rounded = n.round();
759 if (rounded - n).abs() > 1.0e-12 {
760 return Err(cov_error_with_detail(
761 &COV_ERROR_NORMALIZATION_INVALID,
762 "value must be an integer",
763 ));
764 }
765 match rounded as i64 {
766 0 => Ok(CovNormalization::Unbiased),
767 1 => Ok(CovNormalization::Biased),
768 other => Err(cov_error_with_detail(
769 &COV_ERROR_NORMALIZATION_INVALID,
770 format!("expected 0 or 1, received {other}"),
771 )),
772 }
773 }
774 Value::Bool(flag) => Ok(if flag {
775 CovNormalization::Biased
776 } else {
777 CovNormalization::Unbiased
778 }),
779 other => Err(cov_error_with_detail(
780 &COV_ERROR_NORMALIZATION_INVALID,
781 format!("value must be numeric, received {other:?}"),
782 )),
783 }
784}
785
786fn should_treat_as_weight(
787 first: &Value,
788 candidate: &Value,
789 normalization_explicit: bool,
790 rows_option: CovRows,
791) -> BuiltinResult<bool> {
792 let (rows_first, cols_first) = value_rows_cols(first)?;
793 let (rows_candidate, cols_candidate) = value_rows_cols(candidate)?;
794
795 let is_vector = rows_candidate == 1
796 || cols_candidate == 1
797 || rows_candidate * cols_candidate == rows_candidate
798 && (rows_candidate == rows_first || cols_candidate == rows_first);
799
800 if !is_vector {
801 return Ok(false);
802 }
803
804 if rows_candidate != rows_first && cols_candidate != rows_first {
805 return Ok(false);
807 }
808
809 if cols_first == 1 && !normalization_explicit && matches!(rows_option, CovRows::All) {
810 return Ok(false);
812 }
813
814 Ok(true)
815}
816
817fn value_rows_cols(value: &Value) -> BuiltinResult<(usize, usize)> {
818 match value {
819 Value::Tensor(tensor) => Ok((tensor.rows(), tensor.cols())),
820 Value::LogicalArray(array) => {
821 if array.shape.len() > 2 {
822 return Err(cov_error_with_detail(
823 &COV_ERROR_INVALID_ARGUMENT,
824 "inputs must be 2-D matrices or vectors",
825 ));
826 }
827 let rows = if array.shape.is_empty() {
828 1
829 } else {
830 array.shape[0]
831 };
832 let cols = if array.shape.len() >= 2 {
833 array.shape[1]
834 } else {
835 1
836 };
837 Ok((rows, cols))
838 }
839 Value::GpuTensor(handle) => {
840 if handle.shape.len() > 2 {
841 return Err(cov_error_with_detail(
842 &COV_ERROR_INVALID_ARGUMENT,
843 "inputs must be 2-D matrices or vectors",
844 ));
845 }
846 let rows = if handle.shape.is_empty() {
847 1
848 } else {
849 handle.shape[0]
850 };
851 let cols = if handle.shape.len() >= 2 {
852 handle.shape[1]
853 } else {
854 1
855 };
856 Ok((rows, cols))
857 }
858 Value::Num(_) | Value::Int(_) | Value::Bool(_) => Ok((1, 1)),
859 other => Err(cov_error_with_detail(
860 &COV_ERROR_INVALID_ARGUMENT,
861 format!("unsupported input type for shape inspection: {other:?}"),
862 )),
863 }
864}
865
866#[derive(Debug, Clone)]
867struct Matrix {
868 data: Vec<f64>,
869 rows: usize,
870 cols: usize,
871}
872
873impl Matrix {
874 fn from_tensor(name: &str, tensor: Tensor) -> BuiltinResult<Self> {
875 if tensor.shape.len() > 2 {
876 return Err(cov_error_with_detail(
877 &COV_ERROR_INVALID_ARGUMENT,
878 format!("{name}: inputs must be 2-D matrices or vectors"),
879 ));
880 }
881 Ok(Self {
882 rows: tensor.rows(),
883 cols: tensor.cols(),
884 data: tensor.data,
885 })
886 }
887
888 #[inline]
889 fn get(&self, row: usize, col: usize) -> f64 {
890 self.data[row + col * self.rows]
891 }
892
893 #[inline]
894 fn column(&self, col: usize) -> &[f64] {
895 let start = col * self.rows;
896 let end = start + self.rows;
897 &self.data[start..end]
898 }
899}
900
901fn combine_tensors(left: Tensor, right: Option<Tensor>) -> BuiltinResult<Matrix> {
902 let mut matrix = Matrix::from_tensor("cov", left)?;
903 if let Some(second) = right {
904 let right_matrix = Matrix::from_tensor("cov", second)?;
905 if matrix.rows != right_matrix.rows {
906 return Err(cov_error(&COV_ERROR_ROWS_MISMATCH));
907 }
908 matrix.cols += right_matrix.cols;
909 matrix
910 .data
911 .extend_from_slice(&right_matrix.data[..right_matrix.rows * right_matrix.cols]);
912 }
913 Ok(matrix)
914}
915
916fn covariance_dense(matrix: &Matrix, weight: &CovWeightSpec) -> BuiltinResult<Tensor> {
917 let cols = matrix.cols;
918 let rows = matrix.rows;
919
920 if cols == 0 {
921 return Tensor::new(Vec::new(), vec![0, 0]).map_err(cov_internal_error);
922 }
923
924 let mut result = vec![f64::NAN; cols * cols];
925
926 match weight {
927 CovWeightSpec::Scalar(normalization) => {
928 let denom = match normalization {
929 CovNormalization::Unbiased => (rows as f64) - 1.0,
930 CovNormalization::Biased => rows as f64,
931 };
932 if denom <= 0.0 {
933 return Tensor::new(result, vec![cols, cols]).map_err(cov_internal_error);
934 }
935
936 let mut means = vec![0.0; cols];
937 for (col, mean_slot) in means.iter_mut().enumerate() {
938 let column = matrix.column(col);
939 let mut sum = 0.0;
940 let mut valid = true;
941 for &value in column {
942 if !value.is_finite() {
943 valid = false;
944 break;
945 }
946 sum += value;
947 }
948 *mean_slot = if valid { sum / (rows as f64) } else { f64::NAN };
949 }
950
951 for i in 0..cols {
952 for j in i..cols {
953 let value = covariance_unweighted_pair(matrix, i, j, &means, denom);
954 set_entry(&mut result, cols, i, j, sanitize_covariance(i == j, value));
955 }
956 }
957 }
958 CovWeightSpec::Vector(weights) => {
959 if weights.len() != rows {
960 return Err(cov_error_with_detail(
961 &COV_ERROR_WEIGHT_VECTOR_LENGTH_MISMATCH,
962 format!("expected {rows} elements"),
963 ));
964 }
965 let sum_w: f64 = weights.iter().sum();
966 if sum_w <= 0.0 {
967 return Tensor::new(result, vec![cols, cols]).map_err(cov_internal_error);
968 }
969 let denom = sum_w - 1.0;
970 if denom <= 0.0 {
971 return Tensor::new(result, vec![cols, cols]).map_err(cov_internal_error);
972 }
973
974 let mut means = vec![0.0; cols];
975 for (col, mean_slot) in means.iter_mut().enumerate() {
976 let column = matrix.column(col);
977 let mut weighted_sum = 0.0;
978 let mut valid = true;
979 for (row, &value) in column.iter().enumerate() {
980 if !value.is_finite() {
981 valid = false;
982 break;
983 }
984 weighted_sum += weights[row] * value;
985 }
986 *mean_slot = if valid {
987 weighted_sum / sum_w
988 } else {
989 f64::NAN
990 };
991 }
992
993 for i in 0..cols {
994 for j in i..cols {
995 let value = covariance_weighted_pair(matrix, i, j, weights, &means, denom);
996 set_entry(&mut result, cols, i, j, sanitize_covariance(i == j, value));
997 }
998 }
999 }
1000 }
1001
1002 Tensor::new(result, vec![cols, cols]).map_err(cov_internal_error)
1003}
1004
1005fn filter_complete_rows(matrix: &Matrix, weight: CovWeightSpec) -> (Matrix, CovWeightSpec) {
1006 if matrix.rows == 0 {
1007 return (
1008 Matrix {
1009 data: Vec::new(),
1010 rows: 0,
1011 cols: matrix.cols,
1012 },
1013 weight,
1014 );
1015 }
1016
1017 let mut valid_rows = Vec::new();
1018 for row in 0..matrix.rows {
1019 let mut is_valid = true;
1020 for col in 0..matrix.cols {
1021 if !matrix.get(row, col).is_finite() {
1022 is_valid = false;
1023 break;
1024 }
1025 }
1026 if is_valid {
1027 valid_rows.push(row);
1028 }
1029 }
1030
1031 if valid_rows.len() == matrix.rows {
1032 return (matrix.clone(), weight);
1034 }
1035
1036 let mut data = Vec::with_capacity(valid_rows.len() * matrix.cols);
1037 for col in 0..matrix.cols {
1038 for &row in &valid_rows {
1039 data.push(matrix.get(row, col));
1040 }
1041 }
1042
1043 let filtered_matrix = Matrix {
1044 data,
1045 rows: valid_rows.len(),
1046 cols: matrix.cols,
1047 };
1048
1049 let filtered_weight = match weight {
1050 CovWeightSpec::Scalar(norm) => CovWeightSpec::Scalar(norm),
1051 CovWeightSpec::Vector(vec) => {
1052 let mut filtered = Vec::with_capacity(valid_rows.len());
1053 for &row in &valid_rows {
1054 filtered.push(vec[row]);
1055 }
1056 CovWeightSpec::Vector(filtered)
1057 }
1058 };
1059
1060 (filtered_matrix, filtered_weight)
1061}
1062
1063fn covariance_pairwise(matrix: &Matrix, weight: &CovWeightSpec) -> BuiltinResult<Tensor> {
1064 let cols = matrix.cols;
1065 if cols == 0 {
1066 return Tensor::new(Vec::new(), vec![0, 0]).map_err(cov_internal_error);
1067 }
1068 let mut result = vec![f64::NAN; cols * cols];
1069 for i in 0..cols {
1070 let variance = covariance_pair(matrix, i, i, weight);
1071 set_entry(&mut result, cols, i, i, sanitize_covariance(true, variance));
1072 for j in (i + 1)..cols {
1073 let value = covariance_pair(matrix, i, j, weight);
1074 set_entry(&mut result, cols, i, j, sanitize_covariance(false, value));
1075 }
1076 }
1077 Tensor::new(result, vec![cols, cols]).map_err(cov_internal_error)
1078}
1079
1080fn covariance_unweighted_pair(
1081 matrix: &Matrix,
1082 lhs: usize,
1083 rhs: usize,
1084 means: &[f64],
1085 denom: f64,
1086) -> f64 {
1087 if !means[lhs].is_finite() || !means[rhs].is_finite() {
1088 return f64::NAN;
1089 }
1090 let mut accumulator = 0.0;
1091 for row in 0..matrix.rows {
1092 let x = matrix.get(row, lhs);
1093 let y = matrix.get(row, rhs);
1094 if !x.is_finite() || !y.is_finite() {
1095 return f64::NAN;
1096 }
1097 accumulator += (x - means[lhs]) * (y - means[rhs]);
1098 }
1099 accumulator / denom
1100}
1101
1102fn covariance_weighted_pair(
1103 matrix: &Matrix,
1104 lhs: usize,
1105 rhs: usize,
1106 weights: &[f64],
1107 means: &[f64],
1108 denom: f64,
1109) -> f64 {
1110 if !means[lhs].is_finite() || !means[rhs].is_finite() {
1111 return f64::NAN;
1112 }
1113 let mut accumulator = 0.0;
1114 for (row, &weight) in weights.iter().enumerate().take(matrix.rows) {
1115 if weight == 0.0 {
1116 continue;
1117 }
1118 let x = matrix.get(row, lhs);
1119 let y = matrix.get(row, rhs);
1120 if !x.is_finite() || !y.is_finite() {
1121 return f64::NAN;
1122 }
1123 accumulator += weight * (x - means[lhs]) * (y - means[rhs]);
1124 }
1125 accumulator / denom
1126}
1127
1128fn covariance_pair(matrix: &Matrix, lhs: usize, rhs: usize, weight: &CovWeightSpec) -> f64 {
1129 match weight {
1130 CovWeightSpec::Scalar(normalization) => {
1131 let mut xs = Vec::new();
1132 let mut ys = Vec::new();
1133 for row in 0..matrix.rows {
1134 let x = matrix.get(row, lhs);
1135 let y = matrix.get(row, rhs);
1136 if x.is_finite() && y.is_finite() {
1137 xs.push(x);
1138 ys.push(y);
1139 }
1140 }
1141 covariance_unweighted_slice(&xs, &ys, *normalization)
1142 }
1143 CovWeightSpec::Vector(weights) => {
1144 let mut xs = Vec::new();
1145 let mut ys = Vec::new();
1146 let mut ws = Vec::new();
1147 for (row, &weight) in weights.iter().enumerate().take(matrix.rows) {
1148 let x = matrix.get(row, lhs);
1149 let y = matrix.get(row, rhs);
1150 if x.is_finite() && y.is_finite() {
1151 xs.push(x);
1152 ys.push(y);
1153 ws.push(weight);
1154 }
1155 }
1156 covariance_weighted_slice(&xs, &ys, &ws)
1157 }
1158 }
1159}
1160
1161fn covariance_unweighted_slice(xs: &[f64], ys: &[f64], normalization: CovNormalization) -> f64 {
1162 if xs.is_empty() || ys.is_empty() {
1163 return f64::NAN;
1164 }
1165 let n = xs.len().min(ys.len());
1166 if n == 0 {
1167 return f64::NAN;
1168 }
1169 let denom = match normalization {
1170 CovNormalization::Unbiased => (n as f64) - 1.0,
1171 CovNormalization::Biased => n as f64,
1172 };
1173 if denom <= 0.0 {
1174 return f64::NAN;
1175 }
1176 let sum_x: f64 = xs.iter().take(n).sum();
1177 let sum_y: f64 = ys.iter().take(n).sum();
1178 let mean_x = sum_x / (n as f64);
1179 let mean_y = sum_y / (n as f64);
1180 let mut accumulator = 0.0;
1181 for idx in 0..n {
1182 accumulator += (xs[idx] - mean_x) * (ys[idx] - mean_y);
1183 }
1184 accumulator / denom
1185}
1186
1187fn covariance_weighted_slice(xs: &[f64], ys: &[f64], weights: &[f64]) -> f64 {
1188 if xs.is_empty() || ys.is_empty() || weights.is_empty() {
1189 return f64::NAN;
1190 }
1191 let n = xs.len().min(ys.len()).min(weights.len());
1192 if n == 0 {
1193 return f64::NAN;
1194 }
1195 let sum_w: f64 = weights.iter().take(n).sum();
1196 if sum_w <= 0.0 {
1197 return f64::NAN;
1198 }
1199 let denom = sum_w - 1.0;
1200 if denom <= 0.0 {
1201 return f64::NAN;
1202 }
1203 let mut mean_x = 0.0;
1204 let mut mean_y = 0.0;
1205 for idx in 0..n {
1206 mean_x += weights[idx] * xs[idx];
1207 mean_y += weights[idx] * ys[idx];
1208 }
1209 mean_x /= sum_w;
1210 mean_y /= sum_w;
1211 let mut accumulator = 0.0;
1212 for idx in 0..n {
1213 accumulator += weights[idx] * (xs[idx] - mean_x) * (ys[idx] - mean_y);
1214 }
1215 accumulator / denom
1216}
1217
1218fn sanitize_covariance(is_diag: bool, value: f64) -> f64 {
1219 if !value.is_finite() {
1220 return value;
1221 }
1222 if is_diag && value < 0.0 && value > -1.0e-12 {
1223 0.0
1224 } else {
1225 value
1226 }
1227}
1228
1229fn set_entry(buffer: &mut [f64], dim: usize, row: usize, col: usize, value: f64) {
1230 let idx = row + col * dim;
1231 buffer[idx] = value;
1232 if row != col {
1233 let symmetrical = col + row * dim;
1234 buffer[symmetrical] = value;
1235 }
1236}
1237
1238#[cfg(test)]
1239pub(crate) mod tests {
1240 use super::*;
1241 use crate::builtins::common::test_support;
1242 use futures::executor::block_on;
1243 use runmat_builtins::{ResolveContext, Tensor, Type};
1244
1245 fn assert_tensor_close(actual: &Tensor, expected: &[f64], tol: f64) {
1246 let dim = (expected.len() as f64).sqrt() as usize;
1247 assert_eq!(actual.shape, vec![dim, dim], "unexpected tensor shape");
1248 for (idx, (&got, &want)) in actual.data.iter().zip(expected.iter()).enumerate() {
1249 if want.is_nan() {
1250 assert!(
1251 got.is_nan(),
1252 "expected NaN at linear index {idx}, found {got}"
1253 );
1254 } else {
1255 assert!(
1256 (got - want).abs() <= tol,
1257 "mismatch at linear index {idx}: got {got}, expected {want}"
1258 );
1259 }
1260 }
1261 }
1262
1263 #[test]
1264 fn cov_type_preserves_column_count() {
1265 let out = cov_type(
1266 &[Type::Tensor {
1267 shape: Some(vec![Some(5), Some(3)]),
1268 }],
1269 &ResolveContext::new(Vec::new()),
1270 );
1271 assert_eq!(
1272 out,
1273 Type::Tensor {
1274 shape: Some(vec![Some(3), Some(3)])
1275 }
1276 );
1277 }
1278
1279 #[test]
1280 fn cov_type_vector_returns_scalar() {
1281 let out = cov_type(
1282 &[Type::Tensor {
1283 shape: Some(vec![Some(1), Some(4)]),
1284 }],
1285 &ResolveContext::new(Vec::new()),
1286 );
1287 assert_eq!(out, Type::Num);
1288 }
1289
1290 #[test]
1291 fn cov_descriptor_signatures_cover_core_forms() {
1292 let labels: Vec<&str> = COV_DESCRIPTOR
1293 .signatures
1294 .iter()
1295 .map(|sig| sig.label)
1296 .collect();
1297 assert!(labels.contains(&"C = cov(X)"));
1298 assert!(labels.contains(&"C = cov(X, normalization)"));
1299 assert!(labels.contains(&"C = cov(X, Y, w, opt)"));
1300 }
1301
1302 #[cfg(feature = "wgpu")]
1303 fn cov_builtin_sync(value: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
1304 block_on(super::cov_builtin(value, rest))
1305 }
1306
1307 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1308 #[test]
1309 fn cov_matrix_basic() {
1310 let tensor = Tensor::new(
1311 vec![
1312 4.0, 4.2, 3.9, 4.3, 4.1, 2.0, 2.1, 2.0, 2.1, 2.2, 0.60, 0.59, 0.58, 0.62, 0.63,
1315 ],
1316 vec![5, 3],
1317 )
1318 .unwrap();
1319 let result = block_on(cov_builtin(Value::Tensor(tensor), Vec::new())).expect("cov");
1320 let tensor = match result {
1321 Value::Tensor(t) => t,
1322 other => panic!("expected tensor result, got {other:?}"),
1323 };
1324 let expected = [
1325 0.0250, 0.0075, 0.00175, 0.0075, 0.0070, 0.00135, 0.00175, 0.00135, 0.00043,
1328 ];
1329 assert_tensor_close(&tensor, &expected, 1.0e-6);
1330 }
1331
1332 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1333 #[test]
1334 fn cov_two_vectors() {
1335 let x = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();
1336 let y = Tensor::new(vec![10.0, 11.0, 9.0, 12.0], vec![4, 1]).unwrap();
1337 let result = block_on(cov_builtin(Value::Tensor(x), vec![Value::Tensor(y)])).expect("cov");
1338 let tensor = match result {
1339 Value::Tensor(t) => t,
1340 other => panic!("expected tensor result, got {other:?}"),
1341 };
1342 let expected = [
1343 1.6666666666666667,
1344 0.6666666666666666, 0.6666666666666666,
1346 1.6666666666666667,
1347 ];
1348 assert_tensor_close(&tensor, &expected, 1.0e-6);
1349 }
1350
1351 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1352 #[test]
1353 fn cov_weighted_vector() {
1354 let tensor = Tensor::new(
1355 vec![
1356 4.0, 4.2, 3.9, 4.3, 4.1, 2.0, 2.1, 2.0, 2.1, 2.2,
1358 ],
1359 vec![5, 2],
1360 )
1361 .unwrap();
1362 let weights = Tensor::new(vec![1.0, 1.0, 1.0, 2.0, 2.0], vec![5, 1]).unwrap();
1363 let result = block_on(cov_builtin(
1364 Value::Tensor(tensor),
1365 vec![Value::Tensor(weights)],
1366 ))
1367 .expect("cov");
1368 let tensor = match result {
1369 Value::Tensor(t) => t,
1370 other => panic!("expected tensor result, got {other:?}"),
1371 };
1372 let expected = [
1373 0.022380952380952376,
1374 0.004999999999999994, 0.004999999999999994,
1376 0.006666666666666678,
1377 ];
1378 assert_tensor_close(&tensor, &expected, 1.0e-6);
1379 }
1380
1381 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1382 #[test]
1383 fn cov_omitrows() {
1384 let tensor = Tensor::new(
1385 vec![
1386 1.0,
1387 3.0,
1388 f64::NAN,
1389 8.0, f64::NAN,
1391 4.0,
1392 6.0,
1393 9.0, 2.0,
1395 5.0,
1396 7.0,
1397 10.0,
1398 ],
1399 vec![4, 3],
1400 )
1401 .unwrap();
1402 let result = block_on(cov_builtin(
1403 Value::Tensor(tensor),
1404 vec![Value::from("omitrows")],
1405 ))
1406 .expect("cov");
1407 let tensor = match result {
1408 Value::Tensor(t) => t,
1409 other => panic!("expected tensor result, got {other:?}"),
1410 };
1411 let expected = [
1412 12.5, 12.5, 12.5, 12.5, 12.5, 12.5, 12.5, 12.5, 12.5,
1415 ];
1416 assert_tensor_close(&tensor, &expected, 1.0e-6);
1417 }
1418
1419 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1420 #[test]
1421 fn cov_partialrows() {
1422 let tensor = Tensor::new(
1423 vec![
1424 1.0,
1425 4.0,
1426 7.0, 2.0,
1428 f64::NAN,
1429 8.0, f64::NAN,
1431 6.0,
1432 9.0,
1433 ],
1434 vec![3, 3],
1435 )
1436 .unwrap();
1437 let result = block_on(cov_builtin(
1438 Value::Tensor(tensor),
1439 vec![Value::from("partialrows")],
1440 ))
1441 .expect("cov");
1442 let tensor = match result {
1443 Value::Tensor(t) => t,
1444 other => panic!("expected tensor result, got {other:?}"),
1445 };
1446 let expected = [
1447 9.0,
1448 18.0,
1449 4.5, 18.0,
1451 18.0,
1452 f64::NAN, 4.5,
1454 f64::NAN,
1455 4.5,
1456 ];
1457 assert_tensor_close(&tensor, &expected, 1.0e-6);
1458 }
1459
1460 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1461 #[test]
1462 fn cov_mismatched_rows_errors() {
1463 let left = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();
1464 let right = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
1465 let err = block_on(cov_builtin(Value::Tensor(left), vec![Value::Tensor(right)]))
1466 .expect_err("expected mismatch error");
1467 assert_eq!(err.identifier(), COV_ERROR_ROWS_MISMATCH.identifier);
1468 }
1469
1470 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1471 #[test]
1472 fn cov_invalid_flag_errors() {
1473 let tensor = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
1474 let err = block_on(cov_builtin(Value::Tensor(tensor), vec![Value::Num(2.5)]))
1475 .expect_err("expected invalid flag error");
1476 assert_eq!(err.identifier(), COV_ERROR_NORMALIZATION_INVALID.identifier);
1477 }
1478
1479 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1480 #[test]
1481 fn cov_weight_vector_length_mismatch_errors() {
1482 let x = Tensor::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![3, 2]).unwrap();
1483 let y = Tensor::new(vec![10.0, 11.0, 12.0], vec![3, 1]).unwrap();
1484 let w = Tensor::new(vec![1.0, 2.0], vec![2, 1]).unwrap();
1485 let err = block_on(cov_builtin(
1486 Value::Tensor(x),
1487 vec![Value::Tensor(y), Value::Tensor(w)],
1488 ))
1489 .expect_err("expected weight length mismatch");
1490 assert_eq!(
1491 err.identifier(),
1492 COV_ERROR_WEIGHT_VECTOR_LENGTH_MISMATCH.identifier
1493 );
1494 }
1495
1496 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1497 #[test]
1498 fn cov_unknown_rows_option_errors() {
1499 let tensor = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
1500 let err = block_on(cov_builtin(
1501 Value::Tensor(tensor),
1502 vec![Value::from("rows"), Value::from("bogus")],
1503 ))
1504 .expect_err("expected unknown rows option error");
1505 assert_eq!(err.identifier(), COV_ERROR_ROWS_OPTION_UNKNOWN.identifier);
1506 }
1507
1508 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1509 #[test]
1510 fn cov_duplicate_normalization_flag_errors() {
1511 let tensor = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
1512 let err = block_on(cov_builtin(
1513 Value::Tensor(tensor),
1514 vec![Value::Num(0.0), Value::Num(1.0)],
1515 ))
1516 .expect_err("expected duplicate normalization flag error");
1517 assert_eq!(
1518 err.identifier(),
1519 COV_ERROR_NORMALIZATION_DUPLICATE.identifier
1520 );
1521 }
1522
1523 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1524 #[test]
1525 fn cov_too_many_array_arguments_errors() {
1526 let x = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
1527 let y = Tensor::new(vec![4.0, 5.0, 6.0], vec![3, 1]).unwrap();
1528 let w = Tensor::new(vec![1.0, 1.0, 1.0], vec![3, 1]).unwrap();
1529 let z = Tensor::new(vec![7.0, 8.0, 9.0], vec![3, 1]).unwrap();
1530 let err = block_on(cov_builtin(
1531 Value::Tensor(x),
1532 vec![Value::Tensor(y), Value::Tensor(w), Value::Tensor(z)],
1533 ))
1534 .expect_err("expected too many array arguments error");
1535 assert_eq!(
1536 err.identifier(),
1537 COV_ERROR_TOO_MANY_ARRAY_ARGUMENTS.identifier
1538 );
1539 }
1540
1541 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1542 #[test]
1543 fn cov_gpu_roundtrip() {
1544 test_support::with_test_provider(|provider| {
1545 let tensor = Tensor::new(
1546 vec![
1547 4.0, 4.2, 3.9, 4.3, 4.1, 2.0, 2.1, 2.0, 2.1, 2.2,
1549 ],
1550 vec![5, 2],
1551 )
1552 .unwrap();
1553 let view = runmat_accelerate_api::HostTensorView {
1554 data: &tensor.data,
1555 shape: &tensor.shape,
1556 };
1557 let handle = provider.upload(&view).expect("upload");
1558 let result = block_on(cov_builtin(Value::GpuTensor(handle), Vec::new())).expect("cov");
1559 let gathered = test_support::gather(result).expect("gather");
1560 let expected = [
1561 0.0250, 0.0075, 0.0075, 0.0070,
1563 ];
1564 assert_tensor_close(&gathered, &expected, 1.0e-6);
1565 });
1566 }
1567
1568 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1569 #[test]
1570 fn cov_gpu_host_weights_return_resident_result() {
1571 test_support::with_test_provider(|provider| {
1572 let tensor = Tensor::new(
1573 vec![
1574 4.0, 4.2, 3.9, 4.3, 4.1, 2.0, 2.1, 2.0, 2.1, 2.2,
1576 ],
1577 vec![5, 2],
1578 )
1579 .unwrap();
1580 let view = runmat_accelerate_api::HostTensorView {
1581 data: &tensor.data,
1582 shape: &tensor.shape,
1583 };
1584 let handle = provider.upload(&view).expect("upload");
1585 let weights = Tensor::new(vec![1.0, 1.0, 1.0, 2.0, 2.0], vec![5, 1]).unwrap();
1586
1587 let result = block_on(cov_builtin(
1588 Value::GpuTensor(handle),
1589 vec![Value::Tensor(weights)],
1590 ))
1591 .expect("weighted cov");
1592 assert!(matches!(result, Value::GpuTensor(_)));
1593 let gathered = test_support::gather(result).expect("gather");
1594 let expected = [
1595 0.022380952380952376,
1596 0.004999999999999994, 0.004999999999999994,
1598 0.006666666666666678,
1599 ];
1600 assert_tensor_close(&gathered, &expected, 1.0e-6);
1601 });
1602 }
1603
1604 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1605 #[test]
1606 fn cov_gpu_resident_weights_return_resident_result() {
1607 test_support::with_test_provider(|provider| {
1608 let tensor = Tensor::new(
1609 vec![
1610 4.0, 4.2, 3.9, 4.3, 4.1, 2.0, 2.1, 2.0, 2.1, 2.2,
1612 ],
1613 vec![5, 2],
1614 )
1615 .unwrap();
1616 let weights = Tensor::new(vec![1.0, 1.0, 1.0, 2.0, 2.0], vec![1, 5]).unwrap();
1617 let data = provider
1618 .upload(&runmat_accelerate_api::HostTensorView {
1619 data: &tensor.data,
1620 shape: &tensor.shape,
1621 })
1622 .expect("upload data");
1623 let weight_handle = provider
1624 .upload(&runmat_accelerate_api::HostTensorView {
1625 data: &weights.data,
1626 shape: &weights.shape,
1627 })
1628 .expect("upload weights");
1629
1630 let result = block_on(cov_builtin(
1631 Value::GpuTensor(data),
1632 vec![Value::GpuTensor(weight_handle)],
1633 ))
1634 .expect("weighted cov");
1635 assert!(matches!(result, Value::GpuTensor(_)));
1636 let gathered = test_support::gather(result).expect("gather");
1637 let expected = [
1638 0.022380952380952376,
1639 0.004999999999999994, 0.004999999999999994,
1641 0.006666666666666678,
1642 ];
1643 assert_tensor_close(&gathered, &expected, 1.0e-6);
1644 });
1645 }
1646
1647 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1648 #[test]
1649 fn cov_gpu_rejects_negative_resident_weights() {
1650 test_support::with_test_provider(|provider| {
1651 let tensor = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
1652 let weights = Tensor::new(vec![1.0, -1.0], vec![2, 1]).unwrap();
1653 let data = provider
1654 .upload(&runmat_accelerate_api::HostTensorView {
1655 data: &tensor.data,
1656 shape: &tensor.shape,
1657 })
1658 .expect("upload data");
1659 let weight_handle = provider
1660 .upload(&runmat_accelerate_api::HostTensorView {
1661 data: &weights.data,
1662 shape: &weights.shape,
1663 })
1664 .expect("upload weights");
1665
1666 let err = block_on(cov_builtin(
1667 Value::GpuTensor(data),
1668 vec![Value::GpuTensor(weight_handle)],
1669 ))
1670 .expect_err("negative weights should fail");
1671 assert_eq!(err.identifier(), COV_ERROR_INVALID_ARGUMENT.identifier);
1672 });
1673 }
1674
1675 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1676 #[test]
1677 #[cfg(feature = "wgpu")]
1678 fn cov_wgpu_matches_cpu() {
1679 let Ok(provider) = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
1680 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
1681 ) else {
1682 return;
1683 };
1684
1685 let tensor = Tensor::new(
1686 vec![
1687 4.0, 4.2, 3.9, 4.3, 4.1, 2.0, 2.1, 2.0, 2.1, 2.2,
1689 ],
1690 vec![5, 2],
1691 )
1692 .unwrap();
1693
1694 let cpu_result =
1695 block_on(cov_builtin(Value::Tensor(tensor.clone()), Vec::new())).expect("cov");
1696 let cpu_tensor = match cpu_result {
1697 Value::Tensor(t) => t,
1698 other => panic!("expected tensor result, got {other:?}"),
1699 };
1700
1701 let view = runmat_accelerate_api::HostTensorView {
1702 data: &tensor.data,
1703 shape: &tensor.shape,
1704 };
1705 let handle = provider.upload(&view).expect("upload");
1706
1707 let gpu_value = cov_builtin_sync(Value::GpuTensor(handle), Vec::new()).expect("cov");
1708 let gathered = test_support::gather(gpu_value).expect("gather");
1709
1710 assert_tensor_close(&gathered, &cpu_tensor.data, 1.0e-6);
1711 }
1712
1713 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1714 #[test]
1715 #[cfg(feature = "wgpu")]
1716 fn cov_wgpu_weighted_matches_cpu_and_stays_resident() {
1717 let Ok(provider) = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
1718 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
1719 ) else {
1720 return;
1721 };
1722
1723 let tensor = Tensor::new(
1724 vec![
1725 4.0, 4.2, 3.9, 4.3, 4.1, 2.0, 2.1, 2.0, 2.1, 2.2,
1727 ],
1728 vec![5, 2],
1729 )
1730 .unwrap();
1731 let weights = Tensor::new(vec![1.0, 1.0, 1.0, 2.0, 2.0], vec![1, 5]).unwrap();
1732
1733 let cpu_result = block_on(cov_builtin(
1734 Value::Tensor(tensor.clone()),
1735 vec![Value::Tensor(weights.clone())],
1736 ))
1737 .expect("cov");
1738 let cpu_tensor = match cpu_result {
1739 Value::Tensor(t) => t,
1740 other => panic!("expected tensor result, got {other:?}"),
1741 };
1742
1743 let data_handle = provider
1744 .upload(&runmat_accelerate_api::HostTensorView {
1745 data: &tensor.data,
1746 shape: &tensor.shape,
1747 })
1748 .expect("upload data");
1749 let weight_handle = provider
1750 .upload(&runmat_accelerate_api::HostTensorView {
1751 data: &weights.data,
1752 shape: &weights.shape,
1753 })
1754 .expect("upload weights");
1755
1756 let gpu_value = cov_builtin_sync(
1757 Value::GpuTensor(data_handle),
1758 vec![Value::GpuTensor(weight_handle)],
1759 )
1760 .expect("weighted cov");
1761 assert!(matches!(gpu_value, Value::GpuTensor(_)));
1762 let gathered = test_support::gather(gpu_value).expect("gather");
1763
1764 assert_tensor_close(&gathered, &cpu_tensor.data, 1.0e-5);
1765 }
1766}