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
67 lines
2.2 KiB
Rust
67 lines
2.2 KiB
Rust
use super::*;
|
|
use burn_tensor::{Distribution, Int, Tensor, backend::Backend};
|
|
use burn_tensor::{IndexingUpdateOp, Tolerance};
|
|
|
|
#[test]
|
|
fn scatter_should_work_with_multiple_workgroups_2d_dim0() {
|
|
same_as_reference_same_shape(0, [256, 32]);
|
|
}
|
|
|
|
#[test]
|
|
fn scatter_should_work_with_multiple_workgroups_2d_dim1() {
|
|
same_as_reference_same_shape(1, [32, 256]);
|
|
}
|
|
|
|
#[test]
|
|
fn scatter_should_work_with_multiple_workgroups_3d_dim0() {
|
|
same_as_reference_same_shape(0, [256, 6, 6]);
|
|
}
|
|
|
|
#[test]
|
|
fn scatter_should_work_with_multiple_workgroups_3d_dim1() {
|
|
same_as_reference_same_shape(1, [6, 256, 6]);
|
|
}
|
|
|
|
#[test]
|
|
fn scatter_should_work_with_multiple_workgroups_3d_dim2() {
|
|
same_as_reference_same_shape(2, [6, 6, 256]);
|
|
}
|
|
|
|
#[test]
|
|
fn scatter_should_work_with_multiple_workgroups_diff_shapes() {
|
|
same_as_reference_diff_shape(1, [32, 128], [32, 1]);
|
|
}
|
|
|
|
fn same_as_reference_diff_shape<const D: usize>(
|
|
dim: usize,
|
|
shape1: [usize; D],
|
|
shape2: [usize; D],
|
|
) {
|
|
let test_device = Default::default();
|
|
TestBackend::seed(&test_device, 0);
|
|
|
|
let tensor = Tensor::<TestBackend, D>::random(shape1, Distribution::Default, &test_device);
|
|
let value = Tensor::<TestBackend, D>::random(shape2, Distribution::Default, &test_device);
|
|
let indices = Tensor::<TestBackend, 1, Int>::random(
|
|
[shape2.iter().product::<usize>()],
|
|
Distribution::Uniform(0., shape2[dim] as f64),
|
|
&test_device,
|
|
)
|
|
.reshape(shape2);
|
|
let ref_device = Default::default();
|
|
let tensor_ref = Tensor::<ReferenceBackend, D>::from_data(tensor.to_data(), &ref_device);
|
|
let value_ref = Tensor::<ReferenceBackend, D>::from_data(value.to_data(), &ref_device);
|
|
let indices_ref = Tensor::<ReferenceBackend, D, Int>::from_data(indices.to_data(), &ref_device);
|
|
|
|
let actual = tensor.scatter(dim, indices, value, IndexingUpdateOp::Add);
|
|
let expected = tensor_ref.scatter(dim, indices_ref, value_ref, IndexingUpdateOp::Add);
|
|
|
|
expected
|
|
.into_data()
|
|
.assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
|
|
}
|
|
|
|
fn same_as_reference_same_shape<const D: usize>(dim: usize, shape: [usize; D]) {
|
|
same_as_reference_diff_shape(dim, shape, shape);
|
|
}
|