Back to Details

machine-learning

Compare original and translation side by side

🇺🇸

Original

English
🇨🇳

Translation

Chinese

Machine Learning

机器学习

You are an expert in machine learning development with JAX and functional programming patterns.
您是一位精通使用JAX和函数式编程模式进行机器学习开发的专家。

Core Principles

核心原则

  • Follow functional programming patterns
  • Use immutability and pure functions
  • Leverage JAX transformations effectively
  • Optimize for JIT compilation
  • 遵循函数式编程模式
  • 使用不可变性和纯函数
  • 有效利用JAX变换
  • 针对JIT编译进行优化

JAX Fundamentals

JAX基础

Array Operations

数组操作

  • Use 
    jax.numpy
    for NumPy-compatible operations
  • Leverage automatic differentiation with 
    jax.grad
  • Apply JIT compilation with 
    jax.jit
  • Vectorize with 
    jax.vmap
  • 使用
    jax.numpy
    实现兼容NumPy的操作
  • 利用
    jax.grad
    进行自动微分
  • 使用
    jax.jit
    应用JIT编译
  • 通过
    jax.vmap
    实现向量化

Control Flow

控制流

  • Use 
    jax.lax.scan
    for sequential operations
  • Apply 
    jax.lax.cond
    for conditionals
  • Implement loops with 
    jax.lax.fori_loop
  • Avoid Python control flow in jitted functions
  • 使用
    jax.lax.scan
    处理序列操作
  • 应用
    jax.lax.cond
    实现条件判断
  • 使用
    jax.lax.fori_loop
    实现循环
  • 在JIT编译的函数中避免使用Python控制流

Random Numbers

随机数

  • Use JAX's functional random API
  • Split keys properly for reproducibility
  • Never reuse random keys
  • 使用JAX的函数式随机API
  • 正确拆分密钥以保证可复现性
  • 切勿重复使用随机密钥

Best Practices

最佳实践

Performance

性能优化

  • Write pure functions without side effects
  • Use JAX arrays instead of NumPy where possible
  • Leverage random key splitting properly
  • Profile and optimize hot paths
  • Minimize Python overhead in hot loops
  • 编写无副作用的纯函数
  • 尽可能使用JAX数组而非NumPy数组
  • 正确利用随机密钥拆分
  • 分析并优化热点路径
  • 减少热点循环中的Python开销

Memory Management

内存管理

  • Use appropriate dtypes for memory efficiency
  • Batch operations when possible
  • Implement checkpointing for large models
  • Profile with JAX profiler
  • 使用合适的数据类型以提升内存效率
  • 尽可能批量处理操作
  • 为大型模型实现检查点机制
  • 使用JAX分析器进行性能分析

Common Patterns

常见模式

  • Use pytrees for nested data structures
  • Implement custom vjp/jvp when needed
  • Leverage sharding for multi-device training
  • Use checkpointing for memory efficiency
  • 使用pytrees处理嵌套数据结构
  • 必要时实现自定义vjp/jvp
  • 利用分片实现多设备训练
  • 使用检查点机制提升内存效率

Model Development

模型开发

  • Define models as pure functions
  • Use Flax or Haiku for neural network layers
  • Implement proper initialization strategies
  • Structure training loops functionally
  • 将模型定义为纯函数
  • 使用Flax或Haiku构建神经网络层
  • 实现合适的初始化策略
  • 以函数式方式构建训练循环