Loading...

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构建神经网络层
实现合适的初始化策略
以函数式方式构建训练循环