Merge pull request #100 from thomaspinder/revamped_init

Revamped init

Author: daniel-dodd

.github/workflows/workflow-master.yml

@@ -9,7 +9,7 @@ on:
 jobs:
   codecov:
     name: Codecov Workflow
-    runs-on: ubuntu-18.04
+    runs-on: ubuntu-22.04
 
     steps:
       - uses: actions/checkout@v1

docs/nbs/barycentres.ipynb

@@ -120,19 +120,19 @@
     "    D = gpx.Dataset(X=x, y=y)\n",
     "    likelihood = gpx.Gaussian(num_datapoints=n)\n",
     "    posterior = gpx.Prior(kernel=gpx.RBF()) * likelihood\n",
-    "    params, trainables, constrainers, unconstrainers = gpx.initialise(posterior)\n",
+    "    params, trainables, constrainers, unconstrainers = gpx.initialise(posterior, key).unpack()\n",
     "    params = gpx.transform(params, unconstrainers)\n",
     "\n",
     "    objective = jax.jit(posterior.marginal_log_likelihood(D, constrainers, negative=True))\n",
     "\n",
     "    opt = ox.adam(learning_rate=0.01)\n",
-    "    learned_params = gpx.fit(\n",
+    "    learned_params, training_history = gpx.fit(\n",
     "        objective=objective,\n",
     "        trainables=trainables,\n",
     "        params=params,\n",
     "        optax_optim=opt,\n",
     "        n_iters=1000,\n",
-    "    )\n",
+    "    ).unpack()\n",
     "    learned_params = gpx.transform(learned_params, constrainers)\n",
     "    return likelihood(posterior(D, learned_params)(xtest), learned_params)\n",
     "\n",

docs/nbs/classification.ipynb

@@ -140,7 +140,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "params, trainable, constrainer, unconstrainer = gpx.initialise(posterior)\n",
+    "parameter_state = gpx.initialise(posterior)\n",
+    "params, trainable, constrainer, unconstrainer = parameter_state.unpack()\n",
     "params = gpx.transform(params, unconstrainer)\n",
     "\n",
     "mll = jax.jit(posterior.marginal_log_likelihood(D, constrainer, negative=True))"
@@ -166,13 +167,13 @@
    "outputs": [],
    "source": [
     "opt = ox.adam(learning_rate=0.01)\n",
-    "unconstrained_params = gpx.fit(\n",
+    "unconstrained_params, training_history = gpx.fit(\n",
     "    mll,\n",
     "    params,\n",
     "    trainable,\n",
     "    opt,\n",
     "    n_iters=500,\n",
-    ")\n",
+    ").unpack()\n",
     "\n",
     "negative_Hessian = jax.jacfwd(jax.jacrev(mll))(unconstrained_params)[\"latent\"][\"latent\"][:,0,:,0]\n",
     "\n",
@@ -585,7 +586,7 @@
    "custom_cell_magics": "kql"
   },
   "kernelspec": {
-   "display_name": "Python 3.8.9 64-bit",
+   "display_name": "Python 3.9.7 ('gpjax')",
    "language": "python",
    "name": "python3"
   },
@@ -599,11 +600,11 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.9"
+   "version": "3.9.7"
   },
   "vscode": {
    "interpreter": {
-    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
+    "hash": "920091140e6b97de16b405af485d142952a229f5dad61a888f46227f5acb94cf"
    }
   }
  },

docs/nbs/collapsed_vi.ipynb

@@ -31,9 +31,7 @@
     "from jax import jit\n",
     "\n",
     "import gpjax as gpx\n",
-    "import tensorflow as tf\n",
     "\n",
-    "tf.random.set_seed(42)\n",
     "key = jr.PRNGKey(123)"
    ]
   },
@@ -188,21 +186,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "params, trainables, constrainers, unconstrainers = gpx.initialise(sgpr)\n",
+    "params, trainables, constrainers, unconstrainers = gpx.initialise(sgpr, key).unpack()\n",
     "\n",
     "loss_fn = jit(sgpr.elbo(D, constrainers, negative=True))\n",
     "\n",
     "optimiser = ox.adam(learning_rate=0.005)\n",
     "\n",
     "params = gpx.transform(params, unconstrainers)\n",
     "\n",
-    "learned_params = gpx.fit(\n",
+    "learned_params, training_history = gpx.fit(\n",
     "    objective = loss_fn,\n",
     "    params = params,\n",
     "    trainables = trainables,\n",
     "    optax_optim = optimiser,\n",
     "    n_iters=2000,\n",
-    ")\n",
+    ").unpack()\n",
     "learned_params = gpx.transform(learned_params, constrainers)"
    ]
   },
@@ -270,7 +268,7 @@
    "outputs": [],
    "source": [
     "full_rank_model = gpx.Prior(kernel = gpx.RBF()) * gpx.Gaussian(num_datapoints=D.n)\n",
-    "fr_params, fr_trainables, fr_constrainers, fr_unconstrainers = gpx.initialise(full_rank_model)\n",
+    "fr_params, fr_trainables, fr_constrainers, fr_unconstrainers = gpx.initialise(full_rank_model, key).unpack()\n",
     "fr_params = gpx.transform(fr_params, fr_unconstrainers)\n",
     "mll = jit(full_rank_model.marginal_log_likelihood(D, fr_constrainers, negative=True))\n",
     "%timeit mll(fr_params).block_until_ready()"

docs/nbs/graph_kernels.ipynb

@@ -113,7 +113,7 @@
     "kernel = gpx.GraphKernel(laplacian=L)\n",
     "f = gpx.Prior(kernel=kernel)\n",
     "\n",
-    "true_params = f.params\n",
+    "true_params = f._initialise_params(key)\n",
     "true_params[\"kernel\"] = {\n",
     "    \"lengthscale\": jnp.array(2.3),\n",
     "    \"variance\": jnp.array(3.2),\n",
@@ -174,21 +174,21 @@
    "source": [
     "likelihood = gpx.Gaussian(num_datapoints=y.shape[0])\n",
     "posterior = f * likelihood\n",
-    "params, trainable, constrainer, unconstrainer = gpx.initialise(posterior)\n",
+    "params, trainable, constrainer, unconstrainer = gpx.initialise(posterior, key).unpack()\n",
     "params = gpx.transform(params, unconstrainer)\n",
     "\n",
     "mll = jit(\n",
     "    posterior.marginal_log_likelihood(train_data=D, transformations=constrainer, negative=True)\n",
     ")\n",
     "\n",
     "opt = ox.adam(learning_rate=0.01)\n",
-    "learned_params = gpx.fit(\n",
+    "learned_params, training_history = gpx.fit(\n",
     "    objective=mll,\n",
     "    params=params,\n",
     "    trainables=trainable,\n",
     "    optax_optim=opt,\n",
     "    n_iters=1000,\n",
-    ")\n",
+    ").unpack()\n",
     "learned_params = gpx.transform(learned_params, constrainer)"
    ]
   },
@@ -297,6 +297,23 @@
    "display_name": "Python 3.9.7 ('gpjax')",
    "language": "python",
    "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "920091140e6b97de16b405af485d142952a229f5dad61a888f46227f5acb94cf"
+   }
   }
  },
  "nbformat": 4,

docs/nbs/haiku.ipynb

@@ -109,11 +109,10 @@
     "\n",
     "    def initialise(self, dummy_x, key):\n",
     "        nn_params = self.network.init(rng=key, x=dummy_x)\n",
-    "        base_kernel_params = self.base_kernel.params\n",
+    "        base_kernel_params = self.base_kernel._initialise_params(key)\n",
     "        self._params = {**nn_params, **base_kernel_params}\n",
     "\n",
-    "    @property\n",
-    "    def params(self):\n",
+    "    def _initialise_params(self, key):\n",
     "        return self._params"
    ]
   },
@@ -176,7 +175,7 @@
     "likelihood = gpx.Gaussian(num_datapoints=D.n)\n",
     "posterior = prior * likelihood\n",
     "\n",
-    "params, trainables, constrainers, unconstrainers = gpx.initialise(posterior)\n",
+    "params, trainables, constrainers, unconstrainers = gpx.initialise(posterior, key).unpack()\n",
     "params = gpx.transform(params, unconstrainers)"
    ]
   },
@@ -217,13 +216,13 @@
     "    ox.adamw(learning_rate=schedule),\n",
     ")\n",
     "\n",
-    "final_params = gpx.fit(\n",
+    "final_params, training_history = gpx.fit(\n",
     "    mll,\n",
     "    params,\n",
     "    trainables,\n",
     "    opt,\n",
     "    n_iters=5000,\n",
-    ")\n",
+    ").unpack()\n",
     "final_params = gpx.transform(final_params, constrainers)"
    ]
   },
@@ -294,6 +293,23 @@
    "display_name": "Python 3.9.7 ('gpjax')",
    "language": "python",
    "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "920091140e6b97de16b405af485d142952a229f5dad61a888f46227f5acb94cf"
+   }
   }
  },
  "nbformat": 4,

docs/nbs/kernels.ipynb

@@ -70,7 +70,7 @@
     "\n",
     "for k, ax in zip(kernels, axes.ravel()):\n",
     "    prior = gpx.Prior(kernel=k)\n",
-    "    params, _, _, _ = gpx.initialise(prior)\n",
+    "    params, _, _, _ = gpx.initialise(prior, key).unpack()\n",
     "    rv = prior(params)(x)\n",
     "    y = rv.sample(sample_shape=10, seed=key)\n",
     "\n",
@@ -117,7 +117,7 @@
    "outputs": [],
    "source": [
     "print(f\"ARD: {slice_kernel.ard}\")\n",
-    "print(f\"Lengthscales: {slice_kernel.params['lengthscale']}\")"
+    "print(f\"Lengthscales: {slice_kernel._initialise_params(key)['lengthscale']}\")"
    ]
   },
   {
@@ -136,7 +136,7 @@
    "outputs": [],
    "source": [
     "x_matrix = jr.normal(key, shape=(50, 5))\n",
-    "K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)\n",
+    "K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel._initialise_params(key))\n",
     "print(K.shape)"
    ]
   },
@@ -160,14 +160,13 @@
    "outputs": [],
    "source": [
     "k1 = gpx.RBF()\n",
-    "k1._params = {\"lengthscale\": jnp.array(1.0), \"variance\": jnp.array(10.0)}\n",
     "k2 = gpx.Polynomial()\n",
     "sum_k = k1 + k2\n",
     "\n",
     "fig, ax = plt.subplots(ncols=3, figsize=(20, 5))\n",
-    "im0 = ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))\n",
-    "im1 = ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))\n",
-    "im2 = ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))\n",
+    "im0 = ax[0].matshow(gpx.kernels.gram(k1, x, k1._initialise_params(key)))\n",
+    "im1 = ax[1].matshow(gpx.kernels.gram(k2, x, k2._initialise_params(key)))\n",
+    "im2 = ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k._initialise_params(key)))\n",
     "\n",
     "fig.colorbar(im0, ax=ax[0])\n",
     "fig.colorbar(im1, ax=ax[1])\n",
@@ -194,10 +193,10 @@
     "prod_k = k1 * k2 * k3\n",
     "\n",
     "fig, ax = plt.subplots(ncols=4, figsize=(20, 5))\n",
-    "im0 = ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))\n",
-    "im1 = ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))\n",
-    "im2 = ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))\n",
-    "im3 = ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))\n",
+    "im0 = ax[0].matshow(gpx.kernels.gram(k1, x, k1._initialise_params(key)))\n",
+    "im1 = ax[1].matshow(gpx.kernels.gram(k2, x, k2._initialise_params(key)))\n",
+    "im2 = ax[2].matshow(gpx.kernels.gram(k3, x, k3._initialise_params(key)))\n",
+    "im3 = ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k._initialise_params(key)))\n",
     "\n",
     "fig.colorbar(im0, ax=ax[0])\n",
     "fig.colorbar(im1, ax=ax[1])\n",
@@ -362,17 +361,17 @@
     "circlular_posterior = gpx.Prior(kernel=PKern) * likelihood\n",
     "\n",
     "# Initialise parameters and corresponding transformations\n",
-    "params, trainable, constrainer, unconstrainer = gpx.initialise(circlular_posterior)\n",
+    "params, trainable, constrainer, unconstrainer = gpx.initialise(circlular_posterior, key).unpack()\n",
     "\n",
     "# Optimise GP's marginal log-likelihood using Adam\n",
     "mll = jit(circlular_posterior.marginal_log_likelihood(D, constrainer, negative=True))\n",
-    "learned_params = gpx.fit(\n",
+    "learned_params, training_history = gpx.fit(\n",
     "    mll,\n",
     "    params,\n",
     "    trainable,\n",
     "    adam(learning_rate=0.05),\n",
     "    n_iters=1000,\n",
-    ")\n",
+    ").unpack()\n",
     "\n",
     "# Untransform learned parameters\n",
     "final_params = gpx.transform(learned_params, constrainer)"