[Feature] Add JAX integration for FlyDSL kernels#257

Open

wenchenvincent wants to merge 9 commits into

wenchenvincent:feat/jax-integration

Commits on Mar 25, 2026

[Feature] Add JAX integration for FlyDSL kernels

wenchenvincent
authored and
Nightly Sweep
committed
[Feature] Add JAX examples for vectorAdd, tiledCopy, and tiledMma

wenchenvincent
authored and
Nightly Sweep
committed
[Test] Add JAX integration tests

wenchenvincent
authored and
Nightly Sweep
committed
[Bugfix] Address PR review comments from Copilot

wenchenvincent
authored and
Nightly Sweep
committed
[Bugfix] Fix jax.jit with JIT cache and address PR review comments

wenchenvincent
authored and
Nightly Sweep
committed
[Bugfix] Include dynamic layout state in JIT cache key

wenchenvincent
authored and
Nightly Sweep
committed
[CI] Add JAX ROCm installation and skip guard for JAX examples

wenchenvincent
authored and
Nightly Sweep
committed
[Misc] Remove unused imports from JAX integration
wenchenvincent
committed
Merge main into feat/jax-integration
wenchenvincent
committed