Switched to KaTeX for math rendering

It's faster and more configurable than MathJax. PiperOrigin-RevId: 251403540 Change-Id: I29d1af207502e5c14488130ca23217cdc54f8707
google-deepmind · Jun 4, 2019 · f1c2a11 · f1c2a11
1 parent 56f22a1
commit f1c2a11
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 17 deletions.
diff --git a/docs/conf.py b/docs/conf.py
@@ -54,8 +54,8 @@
     'sphinx.ext.doctest',
     'sphinx.ext.inheritance_diagram',
     'sphinx.ext.linkcode',
-    'sphinx.ext.mathjax',
     'sphinx.ext.napoleon',
+    'sphinxcontrib.katex',
 ]
 
 # Add any paths that contain templates here, relative to this directory.

diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -1,2 +1,3 @@
 sphinx>=2.0.1
 sphinx_rtd_theme>=0.4.3
+sphinxcontrib-katex>=0.4.1
diff --git a/sonnet/src/recurrent.py b/sonnet/src/recurrent.py
@@ -712,14 +712,14 @@ class LSTM(RNNCore):
 
   .. math::
 
-     \begin{align*}
-     i_t &= \sigma(W_{ii} x_t + W_{hi} h_{t-1} + b_i) \\
-     f_t &= \sigma(W_{if} x_t + W_{hf} h_{t-1} + b_f) \\
-     g_t &= \tanh(W_{ig} x_t + W_{hg} h_{t-1} + b_g) \\
-     o_t &= \sigma(W_{io} x_t + W_{ho} h_{t-1} + b_o) \\
-     c_t &= f_t c_{t-1} + i_t g_t \\
-     h_t &= o_t \tanh(c_t)
-     \end{align*}
+     \begin{array}{ll}
+     i_t = \sigma(W_{ii} x_t + W_{hi} h_{t-1} + b_i) \\
+     f_t = \sigma(W_{if} x_t + W_{hf} h_{t-1} + b_f) \\
+     g_t = \tanh(W_{ig} x_t + W_{hg} h_{t-1} + b_g) \\
+     o_t = \sigma(W_{io} x_t + W_{ho} h_{t-1} + b_o) \\
+     c_t = f_t c_{t-1} + i_t g_t \\
+     h_t = o_t \tanh(c_t)
+     \end{array}
 
   Where :math:`i_t`, :math:`f_t`, :math:`o_t` are input, forget and
   output gate activations, and :math:`g_t` is a vector of cell updates.
@@ -1050,12 +1050,12 @@ def lstm_with_recurrent_dropout(
 
   .. math::
 
-     \begin{align*}
-     i_t &= \sigma(W_{ii} x_t + W_{hi} d(h_{t-1}) + b_i) \\
-     f_t &= \sigma(W_{if} x_t + W_{hf} d(h_{t-1}) + b_f) \\
-     g_t &= \tanh(W_{ig} x_t + W_{hg} d(h_{t-1}) + b_g) \\
-     o_t &= \sigma(W_{io} x_t + W_{ho} d(h_{t-1}) + b_o)
-     \end{align*}
+     \begin{array}{ll}
+     i_t = \sigma(W_{ii} x_t + W_{hi} d(h_{t-1}) + b_i) \\
+     f_t = \sigma(W_{if} x_t + W_{hf} d(h_{t-1}) + b_f) \\
+     g_t = \tanh(W_{ig} x_t + W_{hg} d(h_{t-1}) + b_g) \\
+     o_t = \sigma(W_{io} x_t + W_{ho} d(h_{t-1}) + b_o)
+     \end{array}
 
   Args:
     hidden_size: Hidden layer size.
@@ -1349,12 +1349,12 @@ class GRU(RNNCore):
 
   .. math::
 
-     \begin{align*}
+     \begin{array}{ll}
      z_t &= \sigma(W_{iz} x_t + W_{hz} h_{t-1} + b_z) \\
      r_t &= \sigma(W_{ir} x_t + W_{hr} h_{t-1} + b_r) \\
      a_t &= \tanh(W_{ia} x_t + W_{ha} (r_t h_{t-1}) + b_a) \\
      h_t &= (1 - z_t) h_{t-1} + z_t a_t
-     \end{align*}
+     \end{array}
 
   where :math:`z_t` and :math:`r_t` are reset and update gates.