MaxCode: Add RAG knowledge base — JAX/Flax docs and targeted migration rules

[gvanica]

Summary

Adds 55 new RAG source documents to the MaxCode knowledge base, significantly expanding the corpus available for retrieval-augmented generation during PyTorch-to-JAX migrations. Also updates the RAG agent and vector DB to support the expanded corpus with targeted retrieval.

Generic sources (22 files)

Reference implementations and documentation that the LLM retrieves as context during conversion:

• JAX/Flax documentation — module API, layers API, setup vs compact, common gotchas, jax.lax primitives
• Flash Linear Attention (FLA) library — gated delta-net layers/models, L2 norm, gated layer norm, rotary embeddings, short convolution, naive delta-rule ops
• Flax attention patterns — example attention, linen attention implementation
• MaxText reference layers — attentions, embeddings, linears, normalizations, Qwen3 and DeepSeek model implementations

Targeted sources (30 files)

Focused migration rules that address specific conversion pitfalls discovered during iterative testing:

Category	Rules
Numerics and dtypes	buffer dtype fidelity, mixed precision, float32 softmax upcast, integer dtype casting
Caching	KV cache prefill/decode, encoder-decoder cache, causal conv1d prefill/decode
Projections and init	fused QKV projection, linear init consistency, weight init patterns, no explicit init for bare layers
Structure	config dataclasses, class hierarchy preservation, default value preservation, source faithfulness
Operations	stop_gradient mapping, triangular masking, WY representation, scan vs for-loop, reduction axis preservation, sum/div vs mean, load balancing loss, MoE capacity routing, cosine similarity
Flax specifics	checkpoint API, train/eval mode, Pallas kernel opportunities, no invented attributes, QKVZ interleaved ordering, dead code helpers

Infrastructure changes

• rag/rag_agent.py — Updated to support targeted RAG retrieval alongside generic retrieval
• rag/vector_db.py — Minor updates for expanded corpus handling

Test plan

• Verify RAG agent loads all 55 new source documents without errors
• Confirm targeted retrieval returns relevant rules for a sample conversion query
• Validate vector DB indexing completes for the expanded corpus

---

Split from #17 — PR 1 of 8

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.