Ben_Kosytorz
3a67c0979c
- Updated stablediffusion crate path from "../stable-diffusion-burn" to "./crates/stable-diffusion-burn" for proper workspace resolution - Enhanced .gitignore to include generated model files (.mpk, .pt, .bin, .safetensors, .ckpt) and user_data directory - Added Cargo.lock to gitignore with appropriate comment - Reorganized IDE files section in gitignore for better clarity - Added newline at end of file for proper formatting
45 lines
1.4 KiB
Rust
45 lines
1.4 KiB
Rust
use super::*;
|
|
use burn_tensor::Tolerance;
|
|
use burn_tensor::{Distribution, Int, Shape, Tensor, backend::Backend};
|
|
|
|
#[test]
|
|
fn gather_should_work_with_multiple_workgroups_dim0() {
|
|
test_same_as_ref([6, 256], 0);
|
|
}
|
|
|
|
#[test]
|
|
fn gather_should_work_with_multiple_workgroups_dim1() {
|
|
test_same_as_ref([6, 256], 1);
|
|
}
|
|
|
|
fn test_same_as_ref<const D: usize>(shape: [usize; D], dim: usize) {
|
|
let device = Default::default();
|
|
TestBackend::seed(&device, 0);
|
|
|
|
let max = shape[dim];
|
|
let shape = Shape::new(shape);
|
|
let tensor =
|
|
Tensor::<TestBackend, D>::random(shape.clone(), Distribution::Default, &Default::default());
|
|
let indices = Tensor::<TestBackend, 1, Int>::from_data(
|
|
Tensor::<TestBackend, 1>::random(
|
|
[shape.num_elements()],
|
|
Distribution::Uniform(0., max as f64),
|
|
&Default::default(),
|
|
)
|
|
.into_data(),
|
|
&Default::default(),
|
|
)
|
|
.reshape(shape);
|
|
let tensor_ref =
|
|
Tensor::<ReferenceBackend, D>::from_data(tensor.to_data(), &Default::default());
|
|
let indices_ref =
|
|
Tensor::<ReferenceBackend, D, Int>::from_data(indices.to_data(), &Default::default());
|
|
|
|
let actual = tensor.gather(dim, indices);
|
|
let expected = tensor_ref.gather(dim, indices_ref);
|
|
|
|
expected
|
|
.into_data()
|
|
.assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
|
|
}
|