docs: document JAX backend for DragonNet

[xrhd]

Proposed changes

Follow-up to #918 addressing the two documentation gaps flagged by @jeongyoonlee during review of the JAX/flax.nnx DragonNet backend.

Types of changes

• Documentation update (non-breaking change which adds documentation)

Checklist

• I have read the CONTRIBUTING doc
• I have signed the CLA
• Lint and unit tests pass locally with my changes
• I have added necessary documentation (if appropriate)

Changes

• docs/examples.rst: add examples/dragonnet_jax_vs_tf to the toctree so Sphinx renders the JAX vs TF benchmark notebook (previously omitted, as noted by @jeongyoonlee).
• docs/installation.rst: document the jax extra alongside tf/torch:
  • intro paragraph now mentions inference.jax
  • pip install causalml[jax] (PyPI)
  • uv add "causalml[jax]" (uv)
  • pip install -e ".[jax]" (Install from source)
  • --runjax added to the Running Tests section
• docs/examples/dragonnet_jax_vs_tf.ipynb: rename the ## Installation markdown heading to ## Setup to avoid an autosectionlabel collision with installation.rst that the new toctree entry surfaces (the label installation was duplicated).

Tests

No code changes (.rst + .ipynb markdown only), so no new tests are required. Verified the documentation build locally:

pip install -r docs/requirements.txt
sphinx-build -b html -D nbsphinx_execute=never docs/ docs/_build/html

python -m http.server 8000 --directory docs/_build/html

Result: build succeeds (exit 0). No new warnings introduced (870 vs 872 baseline warnings — the reduction comes from the heading rename removing a duplicate-label collision). The dragonnet_jax_vs_tf notebook now renders and is linked from the Examples page.

Dependencies

None. No new dependencies are added; the jax extra already exists in pyproject.toml from #918.

References

• Predecessor PR: Add JAX/flax.nnx backend for DragonNet #918 (Add JAX/flax.nnx backend for DragonNet)
• Review comment by @jeongyoonlee requesting these two doc follow-ups (toctree entry + pip install causalml[jax]).
• Agreed plan: handle in a separate follow-up PR dedicated to documentation.

[jeongyoonlee]

LGTM. Thanks!