Skip to content

Commit

Permalink
reduce tensor sizes for 'slow' tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@swfsql
swfsql committed Mar 1, 2024
1 parent a1b6c1a commit 38b20e5
Show file tree
Hide file tree
Showing 5 changed files with 35 additions and 35 deletions.
22 changes: 11 additions & 11 deletions dfdx-core/src/tensor_ops/conv2d/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -218,22 +218,22 @@ fn test_conv2d_s4p3k2() {
#[test]
fn test_batched_conv2d() {
let dev: TestDevice = Default::default();
let x: Tensor<Rank3<3, 28, 28>, TestDtype, _> = dev.sample_normal();
let x: Tensor<Rank3<3, 7, 7>, TestDtype, _> = dev.sample_normal();
let w: Tensor<Rank4<5, 3, 6, 6>, TestDtype, _> = dev.sample_normal();

let y: Tensor<Rank3<5, 9, 9>, _, _, _> =
let y: Tensor<Rank3<5, 2, 2>, _, _, _> =
(x.leaky_trace(), w.clone()).conv2d(Const::<3>, Const::<2>, Const::<1>, Const::<1>);
let y0 = y.retaped::<NoneTape>();
let grads0 = y.square().mean().backward();
let x0 = grads0.get(&x);
let w0 = grads0.get(&w);

let x = x
.broadcast::<Rank4<10, 3, 28, 28>, _>()
.reshape::<Rank4<10, 3, 28, 28>>();
.broadcast::<Rank4<10, 3, 7, 7>, _>()
.reshape::<Rank4<10, 3, 7, 7>>();
assert_eq!(x.strides, x.shape.strides());

let y: Tensor<Rank4<10, 5, 9, 9>, _, _, _> =
let y: Tensor<Rank4<10, 5, 2, 2>, _, _, _> =
(x.leaky_trace(), w.clone()).conv2d(Const::<3>, Const::<2>, Const::<1>, Const::<1>);
for i in 0..10 {
assert_close_to_tensor!(y0, y.retaped::<NoneTape>().select(dev.tensor(i)));
Expand All @@ -245,7 +245,7 @@ fn test_batched_conv2d() {

let x_grad = grads.get(&x) * 10.0;
for i in 0..10 {
assert_close_to_tensor!(x0, x_grad.clone().select(dev.tensor(i)));
assert_close_to_tensor!(x0, x_grad.clone().select(dev.tensor(i)), 3e-6);
}
}

Expand Down Expand Up @@ -405,7 +405,7 @@ fn test_conv2d_grouped() {
fn test_conv2d_grouped_slices() {
const NUM_GROUPS: usize = 3;
let dev: TestDevice = Default::default();
let x: Tensor<Rank4<2, 9, 14, 14>, TestDtype, _> = dev.sample_normal();
let x: Tensor<Rank4<2, 9, 3, 3>, TestDtype, _> = dev.sample_normal();
let w: Tensor<Rank4<15, 3, 3, 3>, TestDtype, _> = dev.sample_normal();

let y = (x.leaky_trace(), w.clone()).conv2d(
Expand All @@ -419,7 +419,7 @@ fn test_conv2d_grouped_slices() {
let x_group = x
.clone()
.slice((.., 3 * i..3 * (i + 1), .., ..))
.realize::<(Const<2>, Const<3>, Const<14>, Const<14>)>();
.realize::<(Const<2>, Const<3>, Const<3>, Const<3>)>();
let w_group = w
.clone()
.slice((5 * i..5 * (i + 1), .., .., ..))
Expand All @@ -428,7 +428,7 @@ fn test_conv2d_grouped_slices() {
let y_group_true = y
.retaped::<NoneTape>()
.slice((.., 5 * i..5 * (i + 1), .., ..))
.realize::<(Const<2>, Const<5>, Const<12>, Const<12>)>();
.realize::<(Const<2>, Const<5>, Const<1>, Const<1>)>();
assert_close_to_tensor!(y_group, y_group_true);
}

Expand All @@ -440,7 +440,7 @@ fn test_conv2d_grouped_slices() {
let x_group = x
.clone()
.slice((.., 3 * i..3 * (i + 1), .., ..))
.realize::<(Const<2>, Const<3>, Const<14>, Const<14>)>();
.realize::<(Const<2>, Const<3>, Const<3>, Const<3>)>();
let w_group = w
.clone()
.slice((5 * i..5 * (i + 1), .., .., ..))
Expand All @@ -452,7 +452,7 @@ fn test_conv2d_grouped_slices() {
let x_grad_group_true = x_grad
.clone()
.slice((.., 3 * i..3 * (i + 1), .., ..))
.realize::<(Const<2>, Const<3>, Const<14>, Const<14>)>();
.realize::<(Const<2>, Const<3>, Const<3>, Const<3>)>();
let w_grad_group_true = w_grad
.clone()
.slice((5 * i..5 * (i + 1), .., .., ..))
Expand Down
10 changes: 5 additions & 5 deletions dfdx-core/src/tensor_ops/convtrans2d/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -280,21 +280,21 @@ fn test_convtrans2d_padded() {
#[test]
fn test_convtrans2d_batched() {
let dev: TestDevice = Default::default();
let x: Tensor<Rank3<3, 28, 28>, TestDtype, _> = dev.sample_normal();
let x: Tensor<Rank3<3, 3, 3>, TestDtype, _> = dev.sample_normal();
let w: Tensor<Rank4<3, 5, 6, 6>, TestDtype, _> = dev.sample_normal();

let y: Tensor<Rank3<5, 83, 83>, _, _, _> =
let y: Tensor<Rank3<5, 8, 8>, _, _, _> =
(x.leaky_trace(), w.clone()).convtrans2d(Const::<3>, Const::<2>, Const::<1>, Const::<1>);
let y0 = y.retaped::<NoneTape>();
let grads0 = y.square().mean().backward();
let x0 = grads0.get(&x);
let w0 = grads0.get(&w);

let x = x
.broadcast::<Rank4<10, 3, 28, 28>, _>()
.reshape::<Rank4<10, 3, 28, 28>>();
.broadcast::<Rank4<10, 3, 3, 3>, _>()
.reshape::<Rank4<10, 3, 3, 3>>();

let y: Tensor<Rank4<10, 5, 83, 83>, _, _, _> =
let y: Tensor<Rank4<10, 5, 8, 8>, _, _, _> =
(x.leaky_trace(), w.clone()).convtrans2d(Const::<3>, Const::<2>, Const::<1>, Const::<1>);
for i in 0..10 {
assert_close_to_tensor!(y0, y.retaped::<NoneTape>().select(dev.tensor(i)), 1e-5);
Expand Down
2 changes: 1 addition & 1 deletion dfdx-core/src/tensor_ops/log_softmax.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ mod tests {
#[test]
fn test_log_softmax_equivalence() {
let dev: TestDevice = Default::default();
let t: Tensor<Rank4<8, 16, 32, 64>, TestDtype, _> = dev.sample_normal();
let t: Tensor<Rank4<2, 3, 2, 3>, TestDtype, _> = dev.sample_normal();
let p = t.leaky_trace().log_softmax::<Axis<3>>();
let p_truth = t.leaky_trace() - t.leaky_trace().logsumexp::<_, Axis<3>>().broadcast();
// we can't create an array as it will overflow the stack
Expand Down
34 changes: 17 additions & 17 deletions dfdx-core/src/tensor_ops/matmul/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -346,21 +346,21 @@ mod tests {
}

{
let a: Tensor<Rank3<10, 5, 3>, TestDtype, _> = dev.zeros();
let a: Tensor<Rank3<6, 4, 3>, TestDtype, _> = dev.zeros();
let b: Tensor<Rank2<3, 2>, TestDtype, _> = dev.zeros();
let _: Tensor<Rank3<10, 5, 2>, TestDtype, _> = a.matmul(b);
let _: Tensor<Rank3<6, 4, 2>, TestDtype, _> = a.matmul(b);
}

{
let a: Tensor<Rank3<10, 5, 3>, TestDtype, _> = dev.zeros();
let b: Tensor<Rank3<10, 3, 2>, TestDtype, _> = dev.zeros();
let _: Tensor<Rank3<10, 5, 2>, TestDtype, _> = a.matmul(b);
let a: Tensor<Rank3<6, 4, 3>, TestDtype, _> = dev.zeros();
let b: Tensor<Rank3<6, 3, 2>, TestDtype, _> = dev.zeros();
let _: Tensor<Rank3<6, 4, 2>, TestDtype, _> = a.matmul(b);
}

{
let a: Tensor<Rank4<10, 20, 5, 3>, TestDtype, _> = dev.zeros();
let b: Tensor<Rank4<10, 20, 3, 2>, TestDtype, _> = dev.zeros();
let _: Tensor<Rank4<10, 20, 5, 2>, TestDtype, _> = a.matmul(b);
let a: Tensor<Rank4<6, 7, 4, 3>, TestDtype, _> = dev.zeros();
let b: Tensor<Rank4<6, 7, 3, 2>, TestDtype, _> = dev.zeros();
let _: Tensor<Rank4<6, 7, 4, 2>, TestDtype, _> = a.matmul(b);
}
}

Expand Down Expand Up @@ -427,7 +427,7 @@ mod tests {

#[test]
fn test_matmul_broadcast() {
const N: usize = 5;
const N: usize = 2;
let dev: TestDevice = Default::default();
let a: Tensor<Rank3<N, 4, 3>, TestDtype, _> = dev.sample_normal();
let a_array = a.array();
Expand Down Expand Up @@ -458,7 +458,7 @@ mod tests {

#[test]
fn test_matmul_broadcast_actual() {
const N: usize = 5;
const N: usize = 2;
let dev: TestDevice = Default::default();
let a: Tensor<Rank3<N, 4, 3>, TestDtype, _> = dev.sample_normal();
let b: Tensor<Rank2<3, 2>, TestDtype, _> = dev.sample_normal();
Expand All @@ -476,9 +476,9 @@ mod tests {
fn test_matmul_batched_3d() {
let dev: TestDevice = Default::default();

let a: Tensor<Rank3<5, 3, 2>, TestDtype, _> = dev.sample_normal();
let a: Tensor<Rank3<2, 3, 2>, TestDtype, _> = dev.sample_normal();
let a_array = a.array();
let b: Tensor<Rank3<5, 2, 4>, TestDtype, _> = dev.sample_normal();
let b: Tensor<Rank3<2, 2, 4>, TestDtype, _> = dev.sample_normal();
let b_array = b.array();
let c = a.leaky_trace().matmul(b.clone());
let c_array = c.array();
Expand All @@ -487,7 +487,7 @@ mod tests {
let g_a = g.get(&a).array();
let g_b = g.get(&b).array();

for i in 0..5 {
for i in 0..2 {
let sub_a = dev.tensor(a_array[i]);
let sub_b = dev.tensor(b_array[i]);
let sub_c = sub_a.leaky_trace().matmul(sub_b.clone());
Expand All @@ -502,9 +502,9 @@ mod tests {
fn test_matmul_batched_4d() {
let dev: TestDevice = Default::default();

let a: Tensor<Rank4<7, 5, 3, 2>, TestDtype, _> = dev.sample_normal();
let a: Tensor<Rank4<2, 3, 3, 2>, TestDtype, _> = dev.sample_normal();
let a_array = a.array();
let b: Tensor<Rank4<7, 5, 2, 4>, TestDtype, _> = dev.sample_normal();
let b: Tensor<Rank4<2, 3, 2, 4>, TestDtype, _> = dev.sample_normal();
let b_array = b.array();
let c = a.leaky_trace().matmul(b.clone());
let c_array = c.array();
Expand All @@ -513,8 +513,8 @@ mod tests {
let g_a = g.get(&a).array();
let g_b = g.get(&b).array();

for i in 0..7 {
for j in 0..5 {
for i in 0..2 {
for j in 0..3 {
let sub_a = dev.tensor(a_array[i][j]);
let sub_b = dev.tensor(b_array[i][j]);
let sub_c = sub_a.leaky_trace().matmul(sub_b.clone());
Expand Down
2 changes: 1 addition & 1 deletion dfdx-core/src/tensor_ops/softmax.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ mod tests {
#[test]
fn test_softmax_equivalence() {
let dev: TestDevice = Default::default();
let t: Tensor<Rank4<8, 16, 32, 64>, TestDtype, _> = dev.sample_normal();
let t: Tensor<Rank4<2, 4, 8, 16>, TestDtype, _> = dev.sample_normal();
let p = t.leaky_trace().softmax::<Axis<3>>();
let p_truth = t.leaky_trace().log_softmax::<Axis<3>>().exp();
// we can't create an array as it will overflow the stack
Expand Down

0 comments on commit 38b20e5

Please sign in to comment.