BEGIN_PUBLIC

Use new lax.cond API (#43) ENND_PUBLIC PiperOrigin-RevId: 320217575
google · Jul 8, 2020 · 156a4f7 · 156a4f7
1 parent 661a4fd
commit 156a4f7
Show file tree

Hide file tree

Showing 11 changed files with 6 additions and 6 deletions.
diff --git a/neural_tangents/predict.py b/neural_tangents/predict.py
@@ -457,9 +457,9 @@ def predict_fn(
     # `fx_test_0`] are evaluated, but also a strictly increasing sequence of
     # timesteps, so we always temporarily append an [almost] `0` at the start.
     identity = lambda x: x
-    t0 = lax.cond(t[0] == 0.,
-                  np.full((1,), -1e-24, t.dtype), identity,
-                  np.zeros((1,), t.dtype), identity)
+    t0 = np.where(t[0] == 0,
+                  np.full((1,), -1e-24, t.dtype),
+                  np.zeros((1,), t.dtype))
     t = np.concatenate([t0, t])
 
     # Solve the ODE.

diff --git a/neural_tangents/utils/utils.py b/neural_tangents/utils/utils.py
@@ -355,10 +355,10 @@ def get_masked_array(x: ArrayOrList,
     if mask_constant is None:
       mask = None
     else:
-      id_fn = lambda m: m
       mask = lax.cond(np.isnan(mask_constant),
-                      np.isnan(x), id_fn,
-                      x == mask_constant, id_fn)
+                      lambda x: np.isnan(x),
+                      lambda x: x == mask_constant,
+                      x)
   else:
     raise TypeError(x, type(x))
 

diff --git a/neural_tangents/tests/batch_test.py → tests/batch_test.py b/neural_tangents/tests/batch_test.py → tests/batch_test.py
diff --git a/neural_tangents/tests/empirical_test.py → tests/empirical_test.py b/neural_tangents/tests/empirical_test.py → tests/empirical_test.py
diff --git a/neural_tangents/tests/function_space_test.py → tests/function_space_test.py b/neural_tangents/tests/function_space_test.py → tests/function_space_test.py
diff --git a/neural_tangents/tests/infinite_fcn_test.py → tests/infinite_fcn_test.py b/neural_tangents/tests/infinite_fcn_test.py → tests/infinite_fcn_test.py
diff --git a/neural_tangents/tests/monte_carlo_test.py → tests/monte_carlo_test.py b/neural_tangents/tests/monte_carlo_test.py → tests/monte_carlo_test.py
diff --git a/neural_tangents/tests/predict_test.py → tests/predict_test.py b/neural_tangents/tests/predict_test.py → tests/predict_test.py
diff --git a/neural_tangents/tests/stax_test.py → tests/stax_test.py b/neural_tangents/tests/stax_test.py → tests/stax_test.py
diff --git a/neural_tangents/tests/utils_test.py → tests/utils_test.py b/neural_tangents/tests/utils_test.py → tests/utils_test.py
diff --git a/neural_tangents/tests/weight_space_test.py → tests/weight_space_test.py b/neural_tangents/tests/weight_space_test.py → tests/weight_space_test.py