Merge pull request #4236 from google:nnx-improve-docs-3

PiperOrigin-RevId: 680950137

docs/.readthedocs.yaml

@@ -0,0 +1,31 @@
+# .readthedocs.yml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.10"
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+  configuration: docs/conf.py
+
+# Optionally build your docs in additional formats such as PDF and ePub
+formats:
+  - htmlzip
+  - epub
+  # - pdf
+
+# Optionally set the version of Python and requirements required to build your docs
+python:
+  install:
+    - method: pip
+      path: .
+      extra_requirements:
+        - all
+        - testing
+        - docs

docs_nnx/nnx_basics.ipynb

@@ -79,9 +79,7 @@
    "source": [
     "`nnx.Variable`'s inner values can be accessed using the `.value` property, however\n",
     "for convenience they implement all numeric operators and can be used directly in\n",
-    "arithmetic expressions (as shown above). Additionally, Variables can passed\n",
-    "to any JAX function as they implement the `__jax_array__` protocol (as long as their\n",
-    "inner value is a JAX array).\n",
+    "arithmetic expressions (as shown above).\n",
     "\n",
     "To actually initialize a Module you simply call the constructor, all the parameters\n",
     "of a Module are usually created eagerly. Since Modules hold their own state methods\n",

docs_nnx/nnx_basics.md

@@ -56,9 +56,7 @@ class Linear(nnx.Module):
 
 `nnx.Variable`'s inner values can be accessed using the `.value` property, however
 for convenience they implement all numeric operators and can be used directly in
-arithmetic expressions (as shown above). Additionally, Variables can passed
-to any JAX function as they implement the `__jax_array__` protocol (as long as their
-inner value is a JAX array).
+arithmetic expressions (as shown above).
 
 To actually initialize a Module you simply call the constructor, all the parameters
 of a Module are usually created eagerly. Since Modules hold their own state methods