use runmat_accelerate_api::{
CorrcoefNormalization, CorrcoefOptions, CorrcoefRows, GpuTensorHandle,
};
use runmat_builtins::{Tensor, Value};
use runmat_macros::runtime_builtin;
use crate::builtins::common::gpu_helpers;
use crate::builtins::common::spec::{
BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
};
use crate::builtins::common::tensor::{self, value_to_string};
use crate::builtins::stats::type_resolvers::corrcoef_type;
use crate::{build_runtime_error, BuiltinResult, RuntimeError};
const NAME: &str = "corrcoef";
fn builtin_error(message: impl Into<String>) -> RuntimeError {
build_runtime_error(message).with_builtin(NAME).build()
}
#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::stats::summary::corrcoef")]
pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
name: "corrcoef",
op_kind: GpuOpKind::Custom("summary-stats"),
supported_precisions: &[ScalarType::F32, ScalarType::F64],
broadcast: BroadcastSemantics::None,
provider_hooks: &[ProviderHook::Custom("corrcoef")],
constant_strategy: ConstantStrategy::InlineLiteral,
residency: ResidencyPolicy::NewHandle,
nan_mode: ReductionNaN::Include,
two_pass_threshold: None,
workgroup_size: None,
accepts_nan_mode: false,
notes: "Uses provider-side corrcoef kernels when rows='all'; other cases fall back to host execution.",
};
#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::stats::summary::corrcoef")]
pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
name: "corrcoef",
shape: ShapeRequirements::Any,
constant_strategy: ConstantStrategy::InlineLiteral,
elementwise: None,
reduction: None,
emits_nan: true,
notes: "Fusion planner treats corrcoef as a non-fusible boundary; GPU execution is provided via a custom provider hook.",
};
#[runtime_builtin(
name = "corrcoef",
category = "stats/summary",
summary = "Compute Pearson correlation coefficients for the columns of matrices or paired data sets.",
keywords = "corrcoef,correlation,statistics,rows,normalization,gpu",
accel = "reduction",
type_resolver(corrcoef_type),
builtin_path = "crate::builtins::stats::summary::corrcoef"
)]
async fn corrcoef_builtin(value: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
let args = CorrcoefArgs::parse(value, rest)?;
if let Some(result) = corrcoef_try_gpu(&args).await? {
return Ok(result);
}
corrcoef_host(args).await
}
pub fn corrcoef_from_tensors(
left: Tensor,
right: Option<Tensor>,
normalization: CorrcoefNormalization,
rows: CorrcoefRows,
) -> BuiltinResult<Tensor> {
let matrix = combine_tensors(left, right)?;
match rows {
CorrcoefRows::All => corrcoef_dense(&matrix, normalization),
CorrcoefRows::Complete => {
let filtered = filter_complete_rows(&matrix);
corrcoef_dense(&filtered, normalization)
}
CorrcoefRows::Pairwise => corrcoef_pairwise(&matrix, normalization),
}
}
#[derive(Debug)]
struct CorrcoefArgs {
first: Value,
second: Option<Value>,
normalization: CorrcoefNormalization,
rows: CorrcoefRows,
}
impl CorrcoefArgs {
fn parse(first: Value, rest: Vec<Value>) -> BuiltinResult<Self> {
let mut second: Option<Value> = None;
let mut normalization: Option<CorrcoefNormalization> = None;
let mut rows = CorrcoefRows::All;
let mut iter = rest.into_iter().peekable();
while let Some(arg) = iter.next() {
match arg {
Value::String(_) | Value::StringArray(_) | Value::CharArray(_) => {
let key = value_to_string(&arg)
.ok_or_else(|| builtin_error("corrcoef: expected string argument"))?
.to_ascii_lowercase();
match key.as_str() {
"rows" => {
let option = iter.next().ok_or_else(|| {
builtin_error("corrcoef: expected a rows option after 'rows'")
})?;
let choice = value_to_string(&option)
.ok_or_else(|| {
builtin_error("corrcoef: rows option must be a string value")
})?
.to_ascii_lowercase();
rows = parse_rows_option(&choice)?;
}
_ => {
return Err(builtin_error(format!("corrcoef: unknown option '{key}'")))
}
}
}
Value::Num(_) | Value::Int(_) | Value::Bool(_) => {
if normalization.is_some() {
return Err(builtin_error(
"corrcoef: normalization flag specified more than once",
));
}
normalization = Some(parse_normalization(arg)?);
}
Value::Tensor(_) | Value::LogicalArray(_) | Value::GpuTensor(_) => {
if second.is_some() {
return Err(builtin_error("corrcoef: too many input arrays"));
}
second = Some(arg);
}
Value::ComplexTensor(_) => {
return Err(builtin_error(
"corrcoef: complex inputs are not supported yet",
));
}
other => {
return Err(builtin_error(format!(
"corrcoef: unsupported argument type {:?}",
other
)))
}
}
}
Ok(Self {
first,
second,
normalization: normalization.unwrap_or(CorrcoefNormalization::Unbiased),
rows,
})
}
}
async fn corrcoef_try_gpu(args: &CorrcoefArgs) -> BuiltinResult<Option<Value>> {
if args.rows != CorrcoefRows::All {
return Ok(None);
}
let provider = match runmat_accelerate_api::provider() {
Some(p) => p,
None => return Ok(None),
};
let first_handle = match &args.first {
Value::GpuTensor(handle) => handle.clone(),
_ => return Ok(None),
};
let maybe_second_handle = match &args.second {
Some(Value::GpuTensor(handle)) => Some(handle.clone()),
Some(_) => return Ok(None),
None => None,
};
let mut owned_concat: Option<GpuTensorHandle> = None;
let matrix_handle = if let Some(second) = maybe_second_handle {
let handles = [first_handle.clone(), second];
match provider.cat(2, &handles) {
Ok(concat) => {
owned_concat = Some(concat.clone());
concat
}
Err(_) => return Ok(None),
}
} else {
first_handle
};
let options = CorrcoefOptions {
normalization: args.normalization,
rows: args.rows,
};
match provider.corrcoef(&matrix_handle, &options).await {
Ok(result) => {
if let Some(temp) = owned_concat {
let _ = provider.free(&temp);
}
Ok(Some(Value::GpuTensor(result)))
}
Err(_) => {
if let Some(temp) = owned_concat {
let _ = provider.free(&temp);
}
Ok(None)
}
}
}
async fn corrcoef_host(args: CorrcoefArgs) -> BuiltinResult<Value> {
let CorrcoefArgs {
first,
second,
normalization,
rows,
} = args;
let left = value_to_tensor_gather(first).await?;
let right = match second {
Some(value) => Some(value_to_tensor_gather(value).await?),
None => None,
};
let tensor = corrcoef_from_tensors(left, right, normalization, rows)?;
Ok(Value::Tensor(tensor))
}
async fn value_to_tensor_gather(value: Value) -> BuiltinResult<Tensor> {
match value {
Value::GpuTensor(handle) => gpu_helpers::gather_tensor_async(&handle).await,
other => tensor::value_into_tensor_for("corrcoef", other).map_err(builtin_error),
}
}
fn parse_rows_option(value: &str) -> BuiltinResult<CorrcoefRows> {
match value {
"all" => Ok(CorrcoefRows::All),
"complete" | "completecase" | "completecases" => Ok(CorrcoefRows::Complete),
"pairwise" | "pairwisecomplete" | "pairwisecompletecase" | "pairwisecompletecases" => {
Ok(CorrcoefRows::Pairwise)
}
other => Err(builtin_error(format!(
"corrcoef: unknown rows option '{other}'"
))),
}
}
fn parse_normalization(value: Value) -> BuiltinResult<CorrcoefNormalization> {
match value {
Value::Int(i) => match i.to_i64() {
0 => Ok(CorrcoefNormalization::Unbiased),
1 => Ok(CorrcoefNormalization::Biased),
other => Err(builtin_error(format!(
"corrcoef: normalization flag must be 0 or 1, received {other}"
))),
},
Value::Num(n) => {
if !n.is_finite() {
return Err(builtin_error("corrcoef: normalization flag must be finite"));
}
let rounded = n.round();
if (rounded - n).abs() > 1.0e-12 {
return Err(builtin_error(
"corrcoef: normalization flag must be an integer",
));
}
match rounded as i64 {
0 => Ok(CorrcoefNormalization::Unbiased),
1 => Ok(CorrcoefNormalization::Biased),
other => Err(builtin_error(format!(
"corrcoef: normalization flag must be 0 or 1, received {other}"
))),
}
}
Value::Bool(b) => Ok(if b {
CorrcoefNormalization::Biased
} else {
CorrcoefNormalization::Unbiased
}),
other => Err(builtin_error(format!(
"corrcoef: normalization flag must be numeric or logical, received {other:?}"
))),
}
}
fn normalization_denominator(norm: CorrcoefNormalization, len: usize) -> f64 {
match norm {
CorrcoefNormalization::Unbiased => (len as f64) - 1.0,
CorrcoefNormalization::Biased => len as f64,
}
}
#[derive(Debug, Clone)]
struct Matrix {
data: Vec<f64>,
rows: usize,
cols: usize,
}
impl Matrix {
fn from_tensor(tensor: Tensor) -> BuiltinResult<Self> {
if tensor.shape.len() > 2 {
return Err(builtin_error(
"corrcoef: inputs must be 2-D matrices or vectors",
));
}
Ok(Self {
rows: tensor.rows(),
cols: tensor.cols(),
data: tensor.data,
})
}
#[inline]
fn get(&self, row: usize, col: usize) -> f64 {
self.data[row + col * self.rows]
}
#[inline]
fn column(&self, col: usize) -> &[f64] {
let start = col * self.rows;
let end = start + self.rows;
&self.data[start..end]
}
}
fn combine_tensors(left: Tensor, right: Option<Tensor>) -> BuiltinResult<Matrix> {
let mut matrix = Matrix::from_tensor(left)?;
if let Some(second) = right {
let right_matrix = Matrix::from_tensor(second)?;
if matrix.rows != right_matrix.rows {
return Err(builtin_error(
"corrcoef: inputs must have the same number of rows",
));
}
matrix.cols += right_matrix.cols;
matrix
.data
.extend_from_slice(&right_matrix.data[..right_matrix.rows * right_matrix.cols]);
}
Ok(matrix)
}
fn filter_complete_rows(matrix: &Matrix) -> Matrix {
if matrix.rows == 0 {
return Matrix {
data: Vec::new(),
rows: 0,
cols: matrix.cols,
};
}
let mut valid_rows = Vec::new();
for row in 0..matrix.rows {
let mut is_valid = true;
for col in 0..matrix.cols {
let value = matrix.get(row, col);
if !value.is_finite() {
is_valid = false;
break;
}
}
if is_valid {
valid_rows.push(row);
}
}
let new_rows = valid_rows.len();
if new_rows == 0 {
return Matrix {
data: Vec::new(),
rows: 0,
cols: matrix.cols,
};
}
let mut data = Vec::with_capacity(new_rows * matrix.cols);
for col in 0..matrix.cols {
for &row in &valid_rows {
data.push(matrix.get(row, col));
}
}
Matrix {
data,
rows: new_rows,
cols: matrix.cols,
}
}
fn corrcoef_dense(matrix: &Matrix, normalization: CorrcoefNormalization) -> BuiltinResult<Tensor> {
let cols = matrix.cols;
if cols == 0 {
return Tensor::new(Vec::new(), vec![0, 0])
.map_err(|e| builtin_error(format!("corrcoef: {e}")));
}
let mut result = vec![f64::NAN; cols * cols];
let rows = matrix.rows;
if rows == 0 {
return Tensor::new(result, vec![cols, cols])
.map_err(|e| builtin_error(format!("corrcoef: {e}")));
}
let denom = normalization_denominator(normalization, rows);
if denom <= 0.0 {
return Tensor::new(result, vec![cols, cols])
.map_err(|e| builtin_error(format!("corrcoef: {e}")));
}
let mut means = vec![0.0; cols];
for (col, mean_slot) in means.iter_mut().enumerate() {
let column = matrix.column(col);
let mut sum = 0.0;
let mut count = 0usize;
for &value in column {
if value.is_finite() {
sum += value;
count += 1;
}
}
*mean_slot = if count > 0 {
sum / (count as f64)
} else {
f64::NAN
};
}
for col in 0..cols {
let mean = means[col];
if !mean.is_finite() {
continue;
}
let mut variance = 0.0;
for row in 0..rows {
let value = matrix.get(row, col);
if !value.is_finite() {
variance = f64::NAN;
break;
}
let dev = value - mean;
variance += dev * dev;
}
if variance.is_nan() {
continue;
}
variance /= denom;
if variance < 0.0 && variance > -1.0e-12 {
variance = 0.0;
}
let stddev = variance.sqrt();
let diag = if stddev > 0.0 { 1.0 } else { f64::NAN };
set_entry(&mut result, cols, col, col, diag);
for other in (col + 1)..cols {
let corr = column_pair_corr(matrix, col, other, &means, denom);
set_entry(&mut result, cols, col, other, corr);
}
}
Tensor::new(result, vec![cols, cols]).map_err(|e| builtin_error(format!("corrcoef: {e}")))
}
fn column_pair_corr(matrix: &Matrix, lhs: usize, rhs: usize, means: &[f64], denom: f64) -> f64 {
let mean_x = means[lhs];
let mean_y = means[rhs];
if !mean_x.is_finite() || !mean_y.is_finite() {
return f64::NAN;
}
let mut var_x = 0.0;
let mut var_y = 0.0;
let mut covariance = 0.0;
for row in 0..matrix.rows {
let a = matrix.get(row, lhs);
let b = matrix.get(row, rhs);
if !a.is_finite() || !b.is_finite() {
return f64::NAN;
}
let dx = a - mean_x;
let dy = b - mean_y;
var_x += dx * dx;
var_y += dy * dy;
covariance += dx * dy;
}
var_x /= denom;
var_y /= denom;
covariance /= denom;
clamp_correlation(divide_covariance(var_x, var_y, covariance))
}
fn corrcoef_pairwise(
matrix: &Matrix,
normalization: CorrcoefNormalization,
) -> BuiltinResult<Tensor> {
let cols = matrix.cols;
if cols == 0 {
return Tensor::new(Vec::new(), vec![0, 0])
.map_err(|e| builtin_error(format!("corrcoef: {e}")));
}
let mut result = vec![f64::NAN; cols * cols];
for col in 0..cols {
set_entry(&mut result, cols, col, col, 1.0);
for other in (col + 1)..cols {
let corr = pairwise_corr(matrix, col, other, normalization);
set_entry(&mut result, cols, col, other, corr);
}
}
Tensor::new(result, vec![cols, cols]).map_err(|e| builtin_error(format!("corrcoef: {e}")))
}
fn pairwise_corr(
matrix: &Matrix,
lhs: usize,
rhs: usize,
normalization: CorrcoefNormalization,
) -> f64 {
let mut xs = Vec::new();
let mut ys = Vec::new();
for row in 0..matrix.rows {
let a = matrix.get(row, lhs);
let b = matrix.get(row, rhs);
if a.is_finite() && b.is_finite() {
xs.push(a);
ys.push(b);
}
}
compute_corr(&xs, &ys, normalization)
}
fn compute_corr(xs: &[f64], ys: &[f64], normalization: CorrcoefNormalization) -> f64 {
if xs.is_empty() || ys.is_empty() {
return f64::NAN;
}
let n = xs.len().min(ys.len());
let denom = normalization_denominator(normalization, n);
if denom <= 0.0 {
return f64::NAN;
}
let sum_x: f64 = xs.iter().take(n).sum();
let sum_y: f64 = ys.iter().take(n).sum();
let mean_x = sum_x / (n as f64);
let mean_y = sum_y / (n as f64);
let mut var_x = 0.0;
let mut var_y = 0.0;
let mut covariance = 0.0;
for i in 0..n {
let dx = xs[i] - mean_x;
let dy = ys[i] - mean_y;
var_x += dx * dx;
var_y += dy * dy;
covariance += dx * dy;
}
var_x /= denom;
var_y /= denom;
covariance /= denom;
if var_x < 0.0 && var_x > -1.0e-12 {
var_x = 0.0;
}
if var_y < 0.0 && var_y > -1.0e-12 {
var_y = 0.0;
}
clamp_correlation(divide_covariance(var_x, var_y, covariance))
}
fn divide_covariance(var_x: f64, var_y: f64, covariance: f64) -> f64 {
if !var_x.is_finite() || !var_y.is_finite() || var_x <= 0.0 || var_y <= 0.0 {
return f64::NAN;
}
let std_x = var_x.sqrt();
let std_y = var_y.sqrt();
if std_x == 0.0 || std_y == 0.0 {
return f64::NAN;
}
covariance / (std_x * std_y)
}
fn clamp_correlation(value: f64) -> f64 {
if value.is_nan() {
return value;
}
if value > 1.0 {
if value - 1.0 < 1.0e-12 {
1.0
} else {
value
}
} else if value < -1.0 {
if -1.0 - value < 1.0e-12 {
-1.0
} else {
value
}
} else {
value
}
}
fn set_entry(buffer: &mut [f64], dim: usize, row: usize, col: usize, value: f64) {
let idx = row + col * dim;
buffer[idx] = value;
if row != col {
let symmetrical = col + row * dim;
buffer[symmetrical] = value;
}
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use crate::builtins::common::test_support;
#[cfg(feature = "wgpu")]
use crate::dispatcher::download_handle_async;
use futures::executor::block_on;
use runmat_builtins::{IntValue, ResolveContext, Tensor, Type, Value};
fn assert_tensor_close(actual: &Tensor, expected: &[f64], tol: f64) {
let dim = (expected.len() as f64).sqrt() as usize;
assert_eq!(actual.shape, vec![dim, dim], "unexpected tensor shape");
for (idx, (&got, &want)) in actual.data.iter().zip(expected.iter()).enumerate() {
if want.is_nan() {
assert!(
got.is_nan(),
"expected NaN at linear index {idx}, found {got}"
);
} else {
assert!(
(got - want).abs() <= tol,
"mismatch at linear index {idx}: got {got}, expected {want}"
);
}
}
}
#[test]
fn corrcoef_type_preserves_column_count() {
let out = corrcoef_type(
&[Type::Tensor {
shape: Some(vec![Some(6), Some(2)]),
}],
&ResolveContext::new(Vec::new()),
);
assert_eq!(
out,
Type::Tensor {
shape: Some(vec![Some(2), Some(2)])
}
);
}
#[test]
fn corrcoef_type_vector_returns_scalar() {
let out = corrcoef_type(
&[Type::Tensor {
shape: Some(vec![Some(4), Some(1)]),
}],
&ResolveContext::new(Vec::new()),
);
assert_eq!(out, Type::Num);
}
fn assert_flow_message(err: RuntimeError, needle: &str) {
assert!(
err.message().contains(needle),
"unexpected error message: {}",
err.message()
);
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn corrcoef_matrix_basic() {
let tensor = Tensor::new(
vec![
1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0, 4.0, 1.0, -1.0, 0.0,
],
vec![4, 3],
)
.unwrap();
let result =
block_on(corrcoef_builtin(Value::Tensor(tensor), Vec::new())).expect("corrcoef");
match result {
Value::Tensor(out) => {
let expected = [
1.0,
1.0,
-0.836_660_026_534,
1.0,
1.0,
-0.836_660_026_534,
-0.836_660_026_534,
-0.836_660_026_534,
1.0,
];
assert_tensor_close(&out, &expected, 1.0e-10);
}
other => panic!("expected tensor result, got {other:?}"),
}
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn corrcoef_two_inputs_matches_concatenation() {
let left = Tensor::new(
vec![
1.0, 2.0, 3.0, 4.0, 4.0, 5.0, 6.0, 7.0,
],
vec![4, 2],
)
.unwrap();
let right = Tensor::new(vec![8.0, 6.0, 7.0, 5.0], vec![4, 1]).unwrap();
let combined = Tensor::new(
vec![
1.0, 2.0, 3.0, 4.0, 4.0, 5.0, 6.0, 7.0, 8.0, 6.0, 7.0, 5.0,
],
vec![4, 3],
)
.unwrap();
let via_two = block_on(corrcoef_builtin(
Value::Tensor(left.clone()),
vec![Value::Tensor(right.clone())],
))
.expect("corrcoef");
let via_combined = block_on(corrcoef_builtin(Value::Tensor(combined), Vec::new()))
.expect("corrcoef combined");
let expected_tensor = match via_combined {
Value::Tensor(t) => t,
_ => panic!("expected tensor output"),
};
let actual_tensor = match via_two {
Value::Tensor(t) => t,
_ => panic!("expected tensor output"),
};
assert_tensor_close(&actual_tensor, &expected_tensor.data, 1.0e-10);
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn corrcoef_rows_complete_ignores_missing() {
let tensor = Tensor::new(
vec![
1.0,
f64::NAN,
3.0,
4.0, 2.0,
5.0,
f64::NAN,
8.0,
],
vec![4, 2],
)
.unwrap();
let result = block_on(corrcoef_builtin(
Value::Tensor(tensor),
vec![Value::from("rows"), Value::from("complete")],
))
.expect("corrcoef");
match result {
Value::Tensor(out) => {
let expected = [
1.0, 1.0, 1.0, 1.0,
];
assert_tensor_close(&out, &expected, 1.0e-10);
}
other => panic!("expected tensor result, got {other:?}"),
}
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn corrcoef_rows_pairwise_staggered_missing() {
let tensor = Tensor::new(
vec![
1.0,
f64::NAN,
4.0,
5.0, 2.0,
5.0,
f64::NAN,
8.0, 3.0,
1.0,
6.0,
f64::NAN,
],
vec![4, 3],
)
.unwrap();
let result = block_on(corrcoef_builtin(
Value::Tensor(tensor),
vec![Value::from("rows"), Value::from("pairwise")],
))
.expect("corrcoef");
match result {
Value::Tensor(out) => {
let expected = [
1.0, 1.0, 1.0, 1.0, 1.0, -1.0, 1.0, -1.0, 1.0,
];
assert_tensor_close(&out, &expected, 1.0e-10);
}
other => panic!("expected tensor result, got {other:?}"),
}
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn corrcoef_flag_one_accepted() {
let tensor = Tensor::new(
vec![
1.0, 3.0, 5.0, 2.0, 4.0, 6.0,
],
vec![3, 2],
)
.unwrap();
let unbiased = block_on(corrcoef_builtin(Value::Tensor(tensor.clone()), Vec::new()))
.expect("unbiased");
let biased = block_on(corrcoef_builtin(
Value::Tensor(tensor),
vec![Value::Int(IntValue::I32(1))],
))
.expect("biased");
let a = match biased {
Value::Tensor(t) => t,
_ => panic!("expected tensor"),
};
let b = match unbiased {
Value::Tensor(t) => t,
_ => panic!("expected tensor"),
};
assert_tensor_close(&a, &b.data, 1.0e-12);
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn corrcoef_gpu_roundtrip() {
test_support::with_test_provider(|provider| {
let tensor = Tensor::new(
vec![
1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0, 4.0, 1.0, -1.0, 0.0,
],
vec![4, 3],
)
.unwrap();
let view = runmat_accelerate_api::HostTensorView {
data: &tensor.data,
shape: &tensor.shape,
};
let handle = provider.upload(&view).expect("upload");
let result =
block_on(corrcoef_builtin(Value::GpuTensor(handle), Vec::new())).expect("corrcoef");
let gathered = test_support::gather(result).expect("gather");
let expected = [
1.0,
1.0,
-0.836_660_026_534,
1.0,
1.0,
-0.836_660_026_534,
-0.836_660_026_534,
-0.836_660_026_534,
1.0,
];
assert_tensor_close(&gathered, &expected, 1.0e-10);
});
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn corrcoef_mismatched_rows_errors() {
let left = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap();
let right = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
let err = block_on(corrcoef_builtin(
Value::Tensor(left),
vec![Value::Tensor(right)],
))
.expect_err("expected mismatch error");
assert_flow_message(err, "same number of rows");
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
fn corrcoef_invalid_flag_errors() {
let tensor = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
let err = block_on(corrcoef_builtin(
Value::Tensor(tensor),
vec![Value::Num(2.5)],
))
.expect_err("expected invalid flag error");
assert_flow_message(err, "normalization flag");
}
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[test]
#[cfg(feature = "wgpu")]
fn corrcoef_wgpu_matches_cpu() {
let _ = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
);
let tensor = Tensor::new(
vec![
1.0, 2.0, 3.0, 4.0, 2.0, 5.0, 6.0, 8.0, 4.0, 1.0, 7.0, 0.0,
],
vec![4, 3],
)
.unwrap();
let cpu = corrcoef_from_tensors(
tensor.clone(),
None,
CorrcoefNormalization::Unbiased,
CorrcoefRows::All,
)
.expect("cpu corrcoef");
let view = runmat_accelerate_api::HostTensorView {
data: &tensor.data,
shape: &tensor.shape,
};
let provider = runmat_accelerate_api::provider().expect("provider");
let handle = provider.upload(&view).expect("upload");
let options = CorrcoefOptions {
normalization: CorrcoefNormalization::Unbiased,
rows: CorrcoefRows::All,
};
let gpu = block_on(provider.corrcoef(&handle, &options)).expect("corrcoef");
let host = block_on(download_handle_async(provider, &gpu)).expect("download");
let gathered =
Tensor::new(host.data.clone(), host.shape.clone()).expect("tensor reconstruction");
assert_tensor_close(&gathered, &cpu.data, 1.0e-6);
}
}