可编程AI与AI DSL可视化实验室
从JAX到LangChain——把AI编程范式的演进以可交互SVG呈现。
01
AI编程范式
声明式与命令式、可微编程、函数式AI与编程抽象层的演进——AI编程如何从手写算子走向可微分的高级抽象。
声明式
描述"做什么"
Keras/Sonnet/Flax高层API
命令式
描述"怎么做"
PyTorch eager模式动态图
可微编程
自动微分原生支持
JAX grad/jit/vmap/pmap变换
函数式
纯函数+不可变
Haiku/PureJAX无副作用
抽象层级
从硬件到应用
硬件→算子→图层→模型→应用
演进趋势
声明+命令融合
PyTorch 2.0 torch.compile
核心张力
灵活性vs性能
动态图易调试vs静态图快执行
02
ML领域特定语言
JAX、PyTorch、TensorFlow三大框架的设计哲学对比：编译优化、自动微分策略与设备映射机制。
JAX
函数式+可组合变换
grad/jit/vmap/pmap正交组合
PyTorch
动态图+eager优先
2.0引入TorchDynamo编译
TensorFlow
静态图+生产优先
tf.function+XLA编译优化
自动微分
反向模式AD
计算图自动构建梯度
编译优化
XLA/TorchInductor
图级融合+内核自动生成
设备映射
SPMD并行
GSPMD/JAX pjit自动分片
生态对比
研究vs生产
PyTorch研究主导/TensorFlow生产强
03
AI应用SDK
LangChain、LlamaIndex、CrewAI等应用层SDK的架构设计：抽象层、设计模式与可组合组件。
LangChain
LLM应用编排框架
Chain/Agent/Memory/Tools
LlamaIndex
数据连接与检索
Index/Query/Retrieval/RAG
CrewAI
多Agent协作
Crew/Agent/Task/Process
设计模式
可组合抽象
Prompt模板/输出解析/重试
抽象层
模型/检索/记忆
统一接口跨提供商
组件复用
模块化设计
工具/检索器/链式组合
LCEL
LangChain表达式
声明式链构建+流式+异步
04
模型编译器
XLA、TVM、MLIR编译器栈：从计算图到优化内核的全流程，图优化、算子融合与内核生成。
XLA
Google编译器
HLO IR→目标代码优化
TVM
Apache张量编译器
Relay IR+AutoTVM自动调优
MLIR
多级中间表示
可扩展编译基础设施
图优化
常量折叠/代数化简
死代码消除/布局优化
算子融合
减少内存往返
逐元素+归约融合策略
内核生成
自动代码生成
Triton/TVM MetaSchedule
自动调优
搜索最优参数
AutoTVM/Ansor/AMPERE
05
AI类型系统
形状推断、张量类型系统、效应系统与类型安全——让AI程序在编译期捕获维度错误与副作用。
形状推断
编译期维度推导
符号形状+约束求解
张量类型
shape+dtype+device
Tensor[D, H, W, C]泛型
效应系统
副作用追踪
IO/Random/Mutable标记
类型安全
编译期防错
维度不匹配在编译期报错
JAX形状
jax.ShapeDtypeStruct
抽象值追踪形状与类型
PyTorch 2.1
FakeTensor模式
无数据张量用于编译分析
未来方向
依赖类型
编译期验证张量维度约束
06
AI运行时
执行引擎、内存管理、调度器与设备抽象——AI模型从编译产物到高效执行的底层基础设施。
执行引擎
图/算子执行
同步/异步/流式执行模式
内存管理
显存分配与复用
内存池/碎片整理/分页
调度器
算子并行调度
DAG调度/流水线/算子重叠
设备抽象
统一设备接口
CPU/GPU/TPU/NPU后端
PagedAttention
vLLM分页KV缓存
减少显存碎片提升吞吐
流式执行
CUDA Stream
计算与传输重叠隐藏延迟
动态调度
运行时自适应
负载均衡+故障恢复
07
AI程序调试
形状调试、梯度检查、性能分析与可视化工具——定位AI程序中的维度错误、梯度异常与性能瓶颈。
形状调试
维度错误定位
逐层形状追踪+断言检查
梯度检查
数值梯度验证
grad_check/有限差分对比
性能分析
瓶颈识别
Chrome Trace/Nsight/torch.profiler
可视化
计算图/注意力/特征
TensorBoard/wandb/Netron
nan/inf检测
数值异常追踪
anomaly_mode自动定位
内存分析
显存占用分解
峰值内存/激活/参数/梯度
分布式调试
多设备一致性
NCCL hang/梯度同步验证
08
AI编程的未来
自然语言编程、AI辅助IDE、程序合成与编程范式变革——当AI成为编程接口，软件工程将如何重塑。
自然语言编程
语言即接口
从代码到意图描述的跃迁
AI辅助IDE
Copilot/Cursor
实时代码生成+重构+调试
程序合成
从规范到代码
输入输出示例→自动生成程序
范式变革
声明式意图编程
描述目标而非步骤
Agent编程
自主代码Agent
Devin/SWE-Agent端到端开发
验证挑战
AI代码可信度
形式化验证+测试生成
人机协作
新型工作流
人类审查+AI生成+自动验证