Update regression files for jax tests

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>
mila-iqia · Nov 28, 2024 · 24f0d3c · 24f0d3c
1 parent e28eedf
commit 24f0d3c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/...test/test_backward_pass_is_reproducible/fashion_mnist_jax_fcnet_jax_image_classifier.yaml b/...test/test_backward_pass_is_reproducible/fashion_mnist_jax_fcnet_jax_image_classifier.yaml
@@ -37,20 +37,20 @@ grads.network.params.1:
 grads.network.params.2:
   device: cuda:0
   max: '1.375e-01'
-  mean: '0.e+00'
+  mean: '-1.490e-09'
   min: '-9.162e-02'
   shape:
   - 10
-  sum: '0.e+00'
+  sum: '-1.490e-08'
 grads.network.params.3:
   device: cuda:0
   max: '3.990e-01'
-  mean: '-1.106e-10'
+  mean: '-3.434e-10'
   min: '-2.054e-01'
   shape:
   - 256
   - 10
-  sum: '-2.831e-07'
+  sum: '-8.792e-07'
 outputs.logits:
   device: cuda:0
   max: '2.656e+00'

diff --git a/...ssifier_test/test_backward_pass_is_reproducible/mnist_jax_fcnet_jax_image_classifier.yaml b/...ssifier_test/test_backward_pass_is_reproducible/mnist_jax_fcnet_jax_image_classifier.yaml
@@ -37,20 +37,20 @@ grads.network.params.1:
 grads.network.params.2:
   device: cuda:0
   max: '4.549e-02'
-  mean: '0.e+00'
+  mean: '-7.451e-10'
   min: '-7.537e-02'
   shape:
   - 10
-  sum: '0.e+00'
+  sum: '-7.451e-09'
 grads.network.params.3:
   device: cuda:0
   max: '7.07e-02'
-  mean: '-5.821e-11'
+  mean: '-4.948e-10'
   min: '-1.064e-01'
   shape:
   - 256
   - 10
-  sum: '-1.490e-07'
+  sum: '-1.267e-06'
 outputs.logits:
   device: cuda:0
   max: '1.85e+00'