Update JAX nn docs

Author: lululxvi

docs/modules/deepxde.nn.jax.rst

@@ -1,8 +1,5 @@
-deepxde.nn.jax package
-======================
-
-Submodules
-----------
+deepxde.nn.jax
+==============
 
 deepxde.nn.jax.fnn module
 -------------------------
@@ -19,11 +16,3 @@ deepxde.nn.jax.nn module
    :members:
    :undoc-members:
    :show-inheritance:
-
-Module contents
----------------
-
-.. automodule:: deepxde.nn.jax
-   :members:
-   :undoc-members:
-   :show-inheritance: