use matten::Tensor;
fn pairwise_euclidean(points: &Tensor) -> Tensor {
let n = points.shape()[0];
let sq = points * points; let _ = sq;
let sq = &(points.clone()) * points;
let row_sq_norms = sq.sum_axis(1);
let gram = points.matmul(&points.transpose());
let col = row_sq_norms.reshape(&[n, 1]);
let row = row_sq_norms.reshape(&[1, n]);
let dist_sq = &(&col + &row) - &(&gram * 2.0);
let clamped: Vec<f64> = dist_sq
.as_slice()
.iter()
.map(|&v| if v < 0.0 { 0.0 } else { v })
.collect();
let clamped_t = Tensor::new(clamped, &[n, n]);
let dists: Vec<f64> = clamped_t.as_slice().iter().map(|v| v.sqrt()).collect();
Tensor::new(dists, &[n, n])
}
fn main() {
let points = Tensor::new(vec![0.0, 0.0, 3.0, 4.0, 6.0, 0.0], &[3, 2]);
let dists = pairwise_euclidean(&points);
println!("pairwise distances:");
for i in 0..3 {
let row = dists.slice().index(i).all().build().unwrap();
println!(" row {i}: {:?}", row.as_slice());
}
assert!((dists.get(&[0, 1]).unwrap() - 5.0).abs() < 1e-9);
assert!((dists.get(&[1, 2]).unwrap() - 5.0).abs() < 1e-9);
println!("Pairwise distances: OK");
}