[DRAFT] Add framework-agnostic operator APIs with JAX support

[mgoldfarb-nvidia]

Summary

• Move framework-neutral operation metadata, validation, result, and GEMM helpers into shared packages while preserving the existing PyTorch wrappers.
• Add lazily imported JAX adapters for RMSNorm, sparse attention, SDPA, and dense, discrete, and grouped GEMM fusions.
• Centralize JAX dtype, layout, divisibility, and pointer-alignment metadata in JaxTensorDesc, deriving CUTLASS TensorSpec values inside JaxApiBase.
• Let XLA own output and workspace allocation; fixed initialization remains descriptor-driven through init_value.
• Preserve CuTe launcher ABI ordering and architecture-specific kernel validation.

Validation

• JAX contracts: 214 passed, 21 skipped, 66 subtests passed.
• Shared and operator contracts: 57 passed, 11 skipped, 104 subtests passed.
• python3 -m compileall -q python/cudnn
• git diff --check

Follow-up

Hardware-dependent tests still need validation on supported NVIDIA GPUs. This is opened as a draft for early API and structure review.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3e139fd4-f696-405a-91f3-5ec529af7624

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}