docs: fix typos and update docstring formatting across multiple modules

optax/_src/alias.py

@@ -2019,7 +2019,7 @@ def rmsprop(
     Objective function: 1.36E+01
 
   References:
-    Hinton, `Overview of mini-batch gradient descent`
+    Hinton, `Overview of mini-batch gradient descent
     <www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf>`_, 2012
 
     Graves, `Generating Sequences With Recurrent Neural Networks

optax/_src/linear_algebra.py

@@ -77,7 +77,7 @@ def power_iteration(
   Args:
     matrix: a square matrix, either as an array or a callable implementing a
       matrix-vector product.
-    v0: initial vector approximating the dominiant eigenvector. If ``matrix`` is
+    v0: initial vector approximating the dominant eigenvector. If ``matrix`` is
       an array of size (n, n), v0 must be a vector of size (n,). If instead
       ``matrix`` is a callable, then v0 must be a tree with the same structure
       as the input of this callable. If this argument is None and ``matrix`` is

optax/contrib/_ademamix.py

@@ -72,7 +72,7 @@ def scale_by_ademamix(
     b1: Exponential decay rate to track the fast EMA.
     b2: Exponential decay rate to track the second moment of past gradients.
     b3: Exponential decay rate to track the slow EMA.
-    alpha: Mixing coefficient in the linear combination for the fast and slow
+    alpha: Mixing coefficient in the linear combination of the fast and slow
       EMAs.
     eps: A small constant applied to denominator outside of the square root (as
       in the Adam paper) to avoid dividing by zero when rescaling.

optax/schedules/_inject.py

@@ -112,7 +112,7 @@ def inject_hyperparams(
 
   Args:
     inner_factory: a function that returns the inner
-      ``optax.GradientTransformation`` with dynamic hyperparameters.
+      :py:class:`.optax.GradientTransformation` with dynamic hyperparameters.
     static_args: a string or iterable of strings specifying which callable
       parameters are not schedules. inject_hyperparams treats all callables as
       schedules by default, so if a hyperparameter is a non-schedule callable,