Mirrors_Framework/deepmind-research
1
0
Fork 0
mirror of https://github.com/google-deepmind/deepmind-research.git synced 2026-08-17 16:43:30 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

JaxTestCase now sets jax_numpy_rank_promotion='raise' by default

Browse Source 
PiperOrigin-RevId: 425896973
This commit is contained in:
Jake VanderPlas
2022-02-16 16:10:24 +00:00
committed by Diego de Las Casas
parent bc869b25d2
commit 35be4a0bb9
2 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
kfac_ferminet_alpha/tests/graph_matcher_test.py
+1
View File
@@ -44,6 +44,7 @@ def tagged_autoencoder(all_params, x_in):
return [[h1, t2], [h2, t2]]
@jtu.with_config(jax_numpy_rank_promotion="allow")
class TestGraphMatcher(jtu.JaxTestCase):
"""Class for running all of the tests for integrating the systems."""
kfac_ferminet_alpha/tests/tracer_test.py
+1
View File
@@ -44,6 +44,7 @@ def autoencoder_aux(all_aux, all_params, x_in):
return [l1, l2 * 0.1], layers_values
@jtu.with_config(jax_numpy_rank_promotion="allow")
class TestTracer(jtu.JaxTestCase):
"""Class for running all of the tests for integrating the systems."""