use crate::{MattenError, Tensor};
fn arg_extreme(
data: &[f64],
operation: &'static str,
want_min: bool,
) -> Result<usize, MattenError> {
if data.iter().any(|v| v.is_nan()) {
return Err(MattenError::InvalidArgument {
operation,
argument: "self",
message: format!("{operation} is undefined for tensors containing NaN"),
});
}
let mut best = 0;
let mut best_val = data[0];
for (i, &v) in data.iter().enumerate().skip(1) {
let better = if want_min { v < best_val } else { v > best_val };
if better {
best = i;
best_val = v;
}
}
Ok(best)
}
impl Tensor {
#[must_use]
pub fn argmin(&self) -> usize {
self.try_argmin().unwrap_or_else(|e| panic!("{e}"))
}
#[must_use]
pub fn argmax(&self) -> usize {
self.try_argmax().unwrap_or_else(|e| panic!("{e}"))
}
pub fn try_argmin(&self) -> Result<usize, MattenError> {
#[cfg(feature = "dynamic")]
if self.is_dynamic() {
return Err(MattenError::Unsupported {
operation: "argmin",
message: "argmin is not supported on dynamic tensors; call try_numeric() first"
.to_string(),
});
}
arg_extreme(&self.data, "argmin", true)
}
pub fn try_argmax(&self) -> Result<usize, MattenError> {
#[cfg(feature = "dynamic")]
if self.is_dynamic() {
return Err(MattenError::Unsupported {
operation: "argmax",
message: "argmax is not supported on dynamic tensors; call try_numeric() first"
.to_string(),
});
}
arg_extreme(&self.data, "argmax", false)
}
}