should mostly fix #665 . Some links fixed in readme and notebooks

README.md

@@ -46,8 +46,8 @@ pip install git+https://github.com/deepmind/optax.git
 
 ## Quickstart
 
-Optax contains implementations of [many popular optimizers](https://optax.readthedocs.io/en/latest/api.html#Common-Optimizers) and
-[loss functions](https://optax.readthedocs.io/en/latest/api.html#common-losses).
+Optax contains implementations of [many popular optimizers](https://optax.readthedocs.io/en/latest/api/common_optimizers.html) and
+[loss functions](https://optax.readthedocs.io/en/latest/api/common_losses.html).
 For example, the following code snippet uses the Adam optimizer from `optax.adam`
 and the mean squared error from `optax.l2_loss`. We initialize the optimizer
 state using the `init` function and `params` of the model.
@@ -212,7 +212,7 @@ critical to support composing a sequence of transformations (e.g. `chain`), as
 well as combining multiple updates to the same parameters (e.g. in multi-task
 settings where different tasks need different sets of gradient transformations).
 
-### Losses ([loss.py](https://github.com/deepmind/optax/blob/master/optax/_src/loss.py))
+### Losses ([loss.py](https://github.com/google-deepmind/optax/tree/master/optax/losses))
 
 Optax provides a number of standard losses used in deep learning, such as
 `l2_loss`, `softmax_cross_entropy`, `cosine_distance`, etc.
@@ -229,14 +229,14 @@ avg_loss = jnp.mean(huber_loss(predictions, targets))
 sum_loss = jnp.sum(huber_loss(predictions, targets))
 ```
 
-### Second Order ([second_order.py](https://github.com/deepmind/optax/blob/master/optax/_src/second_order.py))
+### Second Order ([second_order.py](https://github.com/google-deepmind/optax/tree/master/optax/second_order))
 
 Computing the Hessian or Fisher information matrices for neural networks is
 typically intractable due to the quadratic memory requirements. Solving for the
 diagonals of these matrices is often a better solution. The library offers
 functions for computing these diagonals with sub-quadratic memory requirements.
 
-### Stochastic gradient estimators ([stochastic_gradient_estimators.py](https://github.com/deepmind/optax/blob/master/optax/_src/stochastic_gradient_estimators.py))
+### Stochastic gradient estimators ([stochastic_gradient_estimators.py](https://github.com/google-deepmind/optax/blob/master/optax/monte_carlo/stochastic_gradient_estimators.py))
 
 Stochastic gradient estimators compute Monte Carlo estimates of gradients of
 the expectation of a function under a distribution with respect to the

docs/api/apply_updates.rst

@@ -1,5 +1,6 @@
 Apply Updates
 =============
+.. currentmodule:: optax
 
 .. autosummary::
     apply_updates

docs/design_docs.md

@@ -4,4 +4,4 @@
 
 This proposal authored by [Dian Wu](https://github.com/wdphy16) lead to
 implementations that may be found in
-[optax.experimental](https://optax.readthedocs.io/en/latest/api.html#experimental).
+[optax.experimental](https://optax.readthedocs.io/en/latest/api/contrib.html?complex-valued-optimization#complex-valued-optimization).

examples/quick_start.ipynb

@@ -37,7 +37,7 @@
         "id": "n7kMS9kyM8vM"
       },
       "source": [
-        "In this example, we begin by setting up a simple linear model and a loss function. You can use any other library, such as [haiku](https://github.com/deepmind/dm-haiku) or [Flax](https://github.com/google/flax) to construct your networks. Here, we keep it simple and write it ourselves. The loss function (L2 Loss) comes from optax's [common loss functions](https://optax.readthedocs.io/en/latest/api.html#common-losses) via `optax.l2_loss()`."
+        "In this example, we begin by setting up a simple linear model and a loss function. You can use any other library, such as [haiku](https://github.com/deepmind/dm-haiku) or [Flax](https://github.com/google/flax) to construct your networks. Here, we keep it simple and write it ourselves. The loss function (L2 Loss) comes from optax's [common loss functions](https://optax.readthedocs.io/en/latest/api/common_losses.html) via `optax.l2_loss()`."
       ]
     },
     {
@@ -91,7 +91,7 @@
       "source": [
         "## Basic usage of Optax\n",
         "\n",
-        "Optax contains implementations of [many popular optimizers](https://optax.readthedocs.io/en/latest/api.html#Common-Optimizers) that can be used very simply. For example, the gradient transform for the Adam optimizer is available at `optax.adam()`. For now, let's start by calling the `GradientTransform` object for Adam the `optimizer`. We then initialize the optimizer state using the `init` function and `params` of the network."
+        "Optax contains implementations of [many popular optimizers](https://optax.readthedocs.io/en/latest/api/common_optimizers.html) that can be used very simply. For example, the gradient transform for the Adam optimizer is available at `optax.adam()`. For now, let's start by calling the `GradientTransform` object for Adam the `optimizer`. We then initialize the optimizer state using the `init` function and `params` of the network."
       ]
     },
     {
@@ -120,7 +120,7 @@
         "\n",
         "Next we write the update loop. The `GradientTransform` object contains an `update` function that takes in the current optimizer state and gradients and returns the `updates` that need to be applied to the parameters: `updates, new_opt_state = optimizer.update(grads, opt_state)`.\n",
         "\n",
-        "Optax comes with a few simple [update rules](https://optax.readthedocs.io/en/latest/api.html#apply-updates) that apply the updates from the gradient transforms to the current parameters to return new ones: `new_params = optax.apply_updates(params, updates)`.\n",
+        "Optax comes with a few simple [update rules](https://optax.readthedocs.io/en/latest/api/apply_updates.html) that apply the updates from the gradient transforms to the current parameters to return new ones: `new_params = optax.apply_updates(params, updates)`.\n",
         "\n"
       ]
     },