update from bleepblop

Author: cottrell

README.md

@@ -93,3 +93,4 @@ The functions $u$, $v$ must be monotonic in $\tau$ and $y_c$ respectively.
 - randomize $\tau$ sampling during training in jax instead of grid ...
 - regularization/calibration/conformal prediction
 - review the translation to jax and make sure arch is actually the same as tf
+- make sure data is pure and deterministic across envs

figs/jax/cdfloss_crps.png

@@ -13,6 +13,7 @@ def get_data(seed=1, m=250, n_x=1, n_tau=11, L=2):
     y ~ N(mu(x), sigma(x))
     """
     random.seed(seed)
+    np.random.seed(seed)
     x = (2 * np.random.rand(m, n_x).astype(np.float64) - 1) * 2
     i = np.argsort(x[:, 0])
     x = x[i]  # to make plotting nicer

figs/jax/cdfloss_logistic.png

@@ -161,7 +161,7 @@ def __init__(self, tau_dims, x_dims, final_dims, *, key):
         # Tau network layers
         tau_keys = jax.random.split(keys[0], len(tau_dims) - 1)
         self.tau_layers = [
-            Dense(
+            NonNegDense(
                 tau_dims[i],
                 tau_dims[i + 1],
                 key=tau_keys[i],
@@ -184,7 +184,8 @@ def __init__(self, tau_dims, x_dims, final_dims, *, key):
                 final_dims[i],
                 final_dims[i + 1],
                 key=final_keys[i],
-                activation=(jax.nn.tanh if i < len(final_dims) - 2 else None),
+                activation=None,
+                # activation=(jax.nn.tanh if i < len(final_dims) - 2 else None),
             )
             for i in range(len(final_dims) - 1)
         ]