diff --git a/neural_tangents/_src/stax/branching.py b/neural_tangents/_src/stax/branching.py
index cc8e1cd5..f9443785 100644
--- a/neural_tangents/_src/stax/branching.py
+++ b/neural_tangents/_src/stax/branching.py
@@ -226,7 +226,7 @@ def kernel_fn(ks: Kernels, **kwargs) -> Kernel:
         raise NotImplementedError(
             '`FanInConcat` layer along the non-channel axis is only implemented'
             'for the case if all input layers guaranteed to be mean-zero '
-            'Gaussian, i.e. having all `is_gaussian set to `True`.')
+            'Gaussian, i.e. having all `is_gaussian` set to `True`.')
     else:
       # TODO(romann): allow nonlinearity after channelwise concatenation.
       # TODO(romann): support concatenating different channelwise masks.
diff --git a/neural_tangents/_src/stax/elementwise.py b/neural_tangents/_src/stax/elementwise.py
index 9f149419..37949bbe 100644
--- a/neural_tangents/_src/stax/elementwise.py
+++ b/neural_tangents/_src/stax/elementwise.py
@@ -912,10 +912,12 @@ def nngp_fn_diag(nngp: np.ndarray) -> np.ndarray:
 @layer
 @supports_masking(remask_kernel=False)
 def Polynomial(coef: Sequence[float]) -> InternalLayer:
-  """Polynomials, i.e. `coef[0] + coef[1] * x + ... + coef[n] * x**n`.
+  """Polynomials, i.e. `coef[0] + coef[1] * x + … + coef[n] * x**n`.
 
   Args:
-    coef: a sequence of coefficients. Follows `numpy.polynomial.Polynomial` API.
+    coef:
+      a sequence of coefficients. Follows
+      :class:`numpy.polynomial.polynomial.Polynomial` API.
 
   Returns:
     `(init_fn, apply_fn, kernel_fn)`.
diff --git a/neural_tangents/_src/stax/linear.py b/neural_tangents/_src/stax/linear.py
index 1aca8e5f..5d918f6a 100644
--- a/neural_tangents/_src/stax/linear.py
+++ b/neural_tangents/_src/stax/linear.py
@@ -3162,7 +3162,7 @@ def _conv_kernel_full_spatial_unshared(
   Used when `kernel.diagonal_spatial == False`. Has the same outputs on the
   spatial diagonal as `_conv_kernel_full_spatial_shared`, but `0` in all
   off-spatial-diagonal entries. The diagonal entries are computed via calling
-  ``_conv_kernel_diagonal_spatial`.
+  `_conv_kernel_diagonal_spatial`.
 
   Args:
     lhs: