use super::*;
#[test]
fn test_fuse_dims_contiguous() {
let dims = [3, 4];
let strides1 = [1isize, 3];
let strides2 = [1isize, 3];
let all_strides: Vec<&[isize]> = vec![&strides1, &strides2];
let fused = fuse_dims(&dims, &all_strides);
assert_eq!(fused, vec![12, 1]);
}
#[test]
fn test_fuse_dims_non_contiguous() {
let dims = [3, 4];
let strides1 = [1isize, 10];
let all_strides: Vec<&[isize]> = vec![&strides1];
let fused = fuse_dims(&dims, &all_strides);
assert_eq!(fused, vec![3, 4]);
}
#[test]
fn test_fuse_dims_bilateral_all_contiguous() {
let dims = vec![2, 2, 2, 2];
let src_strides = vec![1, 2, 4, 8];
let dst_strides = vec![1, 2, 4, 8];
let (fd, fs, fds) = fuse_dims_bilateral(&dims, &src_strides, &dst_strides);
assert_eq!(&fd[..], &[16]);
assert_eq!(&fs[..], &[1]);
assert_eq!(&fds[..], &[1]);
}
#[test]
fn test_fuse_dims_bilateral_partial() {
let dims = vec![2, 3, 4];
let src_strides = vec![1, 2, 100]; let dst_strides = vec![1, 2, 6]; let (fd, fs, fds) = fuse_dims_bilateral(&dims, &src_strides, &dst_strides);
assert_eq!(&fd[..], &[6, 4]); assert_eq!(&fs[..], &[1, 100]);
assert_eq!(&fds[..], &[1, 6]);
}
#[test]
fn test_fuse_dims_bilateral_scattered() {
let dims = vec![2, 2, 2];
let src_strides = vec![1, 4194304, 2]; let dst_strides = vec![1, 2, 4]; let (fd, fs, fds) = fuse_dims_bilateral(&dims, &src_strides, &dst_strides);
assert_eq!(&fd[..], &[2, 2, 2]); assert_eq!(&fs[..], &[1, 4194304, 2]);
assert_eq!(&fds[..], &[1, 2, 4]);
}
#[test]
fn test_compress_dims_removes_fused() {
let dims = vec![12usize, 1];
let strides = vec![vec![1isize, 3]];
let (cd, cs) = compress_dims(&dims, &strides);
assert_eq!(cd, vec![12]);
assert_eq!(cs, vec![vec![1]]);
}
#[test]
fn test_compute_costs() {
let strides1 = [1isize, 4, 0];
let strides2 = [2isize, 1, 0];
let all_strides: Vec<&[isize]> = vec![&strides1, &strides2];
let costs = compute_costs(&all_strides);
assert_eq!(costs, vec![2, 2, 1]);
}
#[test]
fn test_compute_importance_uses_each_array_row() {
let dims = [2, 3];
let strides0 = [1isize, 2];
let strides1 = [2isize, 1];
let all_strides: Vec<&[isize]> = vec![&strides0, &strides1];
let index_orders = vec![vec![0usize, 1], vec![1, 0]];
assert_eq!(
compute_importance(&dims, &all_strides, &index_orders),
vec![36, 24]
);
}
#[test]
fn test_compute_importance_accumulates_every_array() {
let dims = [2, 2];
let strides = [1isize, 2];
let two_strides: Vec<&[isize]> = vec![&strides, &strides];
let three_strides: Vec<&[isize]> = vec![&strides, &strides, &strides];
let with_two = vec![vec![0usize, 1], vec![0, 1]];
let with_three = vec![vec![0usize, 1], vec![0, 1], vec![0, 1]];
assert_eq!(
compute_importance(&dims, &two_strides, &with_two),
vec![48, 12]
);
assert_eq!(
compute_importance(&dims, &three_strides, &with_three),
vec![256, 32]
);
}
#[test]
fn test_compute_importance_short_row_contributes_its_prefix() {
let dims = [2, 2];
let full = [1isize, 2];
let broadcast = [0isize];
let all_strides: Vec<&[isize]> = vec![&full, &full, &broadcast];
let index_orders = vec![vec![1usize, 2], vec![1, 2], vec![1]];
assert_eq!(
compute_importance(&dims, &all_strides, &index_orders),
vec![32, 3]
);
}
#[test]
fn test_compute_importance_long_row_ignores_its_suffix() {
let dims = [2, 2];
let full = [1isize, 2];
let all_strides: Vec<&[isize]> = vec![&full, &full, &full];
let index_orders = vec![vec![1usize, 2], vec![1, 2], vec![1, 2, 1]];
assert_eq!(
compute_importance(&dims, &all_strides, &index_orders),
vec![32, 4]
);
}