实践：实现有状态函数

在JAX中管理状态时，主要思想是函数纯粹性：函数不应修改外部状态或有隐藏的副作用。相反，它们应将当前状态作为输入，并返回新状态作为输出，同时包含任何结果。这种“状态输入，状态输出”模式在使用JAX变换（如jit和grad）时非常重要。

这些练习将帮助您熟悉在不同情况下实现这种模式。我们将从简单的开始，逐步完成一个类似于机器学习优化组成部分的任务。请记住，JAX经常使用PyTrees（嵌套元组、列表、字典）来方便地处理复杂状态。

练习1：使用JIT实现有状态计数器

让我们回顾一下有状态计数器的例子。您的任务是实现一个update_counter函数，它接受当前计数（状态）和增量值。它应返回新计数以及增量值本身作为辅助输出。然后，对此函数应用jax.jit并进行测试。

说明：

1. 导入JAX和jax.numpy。
2. 定义update_counter函数。它应接受count和increment作为参数。
3. 在函数内部，计算new_count = count + 1。
4. 函数应返回一个元组：(new_count, increment)。第一个元素是更新后的状态，第二个是辅助输出。
5. 使用jax.jit创建此函数的jit编译版本。
6. 初始化计数器状态（例如，count = 0）。
7. 使用Python for循环多次调用jit编译的函数（例如，5次）。在每次迭代中：
   • 传递当前count和一个示例increment值（例如，循环索引）。
   • 将返回的元组解包为新的count（更新下一次迭代的状态）和returned_increment。
   • 打印更新后的计数和返回的增量值。

import jax
import jax.numpy as jnp

# 1. 定义有状态函数
def update_counter(count, increment):
  """递增计数并返回新计数和增量值。"""
  new_count = count + 1
  # 返回 (新状态, 辅助输出)
  return new_count, increment

# 2. JIT编译函数
jitted_update_counter = jax.jit(update_counter)

# 3. 初始化状态
current_count = 0
print(f"初始计数：{current_count}")

# 4. 运行循环，每次更新状态
num_steps = 5
for i in range(num_steps):
  # 传入当前状态，获取新状态和输出
  current_count, returned_increment = jitted_update_counter(current_count, i)
  print(f"步骤 {i+1}: 新计数 = {current_count}, 返回增量 = {returned_increment}")

预期输出：

初始计数：0
步骤 1: 新计数 = 1, 返回增量 = 0
步骤 2: 新计数 = 2, 返回增量 = 1
步骤 3: 新计数 = 3, 返回增量 = 2
步骤 4: 新计数 = 4, 返回增量 = 3
步骤 5: 新计数 = 5, 返回增量 = 4

这个简单例子展示了主要模式：循环在编译函数外部管理状态（current_count），在每一步中将其传入并接收更新后的版本。

练习2：计算简单移动平均

现在，让我们实现一个函数来计算简单移动平均。移动平均是根据新的输入值更新的。我们还需要记录目前为止已见值的总和以及值的总数量。

说明：

1. 定义一个update_moving_average函数，它接受state和new_value作为输入。
2. state将是一个元组(current_sum, count)。适当地初始化它（例如，(0.0, 0)）。
3. 在函数内部：
   • 解包state元组。
   • 计算new_sum = current_sum + new_value。
   • 计算new_count = count + 1。
   • 计算current_average = new_sum / new_count。如有必要，处理count可能为0的初始情况（尽管此处加1可避免除以零）。
   • 打包更新后的状态：new_state = (new_sum, new_count)。
4. 函数应返回(new_state, current_average)。
5. 创建一个数字序列示例（例如，一个JAX数组）。
6. 初始化状态。
7. 迭代遍历序列，为每个数字调用update_moving_average。在每次迭代中更新状态变量并打印计算出的平均值。
8. 可选： 对update_moving_average应用jax.jit并查看它是否正常工作。

import jax
import jax.numpy as jnp

# 1. 定义用于移动平均的有状态函数
def update_moving_average(state, new_value):
  """使用新值更新移动平均状态。"""
  current_sum, count = state # 解包状态

  new_sum = current_sum + new_value
  new_count = count + 1
  current_average = new_sum / new_count

  new_state = (new_sum, new_count) # 打包新状态
  return new_state, current_average

# 可选：JIT编译函数
# jitted_update_moving_average = jax.jit(update_moving_average)
# 如果您取消注释此行，请在下方使用 jitted_update_moving_average

# 2. 示例数据和初始状态
data_sequence = jnp.array([2.0, 4.0, 6.0, 8.0, 10.0])
initial_state = (0.0, 0) # (总和, 计数)
print(f"初始状态 (总和, 计数): {initial_state}")
print(f"数据序列: {data_sequence}")

# 3. 迭代并更新
current_state = initial_state
for i, value in enumerate(data_sequence):
  # 传入当前状态和值，获取新状态和平均值
  current_state, avg = update_moving_average(current_state, value)
  # 如果正在JIT编译，请使用此行替代：
  # current_state, avg = jitted_update_moving_average(current_state, value)
  print(f"值 {value:.1f} 后: 新状态 = ({current_state[0]:.1f}, {current_state[1]}), 移动平均 = {avg:.2f}")

预期输出：

初始状态 (总和, 计数): (0.0, 0)
数据序列: [ 2.  4.  6.  8. 10.]
值 2.0 后: 新状态 = (2.0, 1), 移动平均 = 2.00
值 4.0 后: 新状态 = (6.0, 2), 移动平均 = 3.00
值 6.0 后: 新状态 = (12.0, 3), 移动平均 = 4.00
值 8.0 后: 新状态 = (20.0, 4), 移动平均 = 5.00
值 10.0 后: 新状态 = (30.0, 5), 移动平均 = 6.00

在这里，状态是一个元组，一个简单的PyTree。函数正确地更新并返回这个结构化状态以及计算出的平均值。JIT编译应该可以工作，因为函数是纯粹的并遵循状态传递模式。

练习3：单步梯度下降

本练习模拟了像梯度下降这样的优化算法中的一个更新步骤。我们将把被优化的参数作为状态进行管理。我们的目标是使一个简单函数最小化，例如 $f(x) = x^2$。

说明：

1. 定义要最小化的损失函数：loss_fn(x) = x**2。
2. 使用jax.grad定义梯度函数：grad_fn = jax.grad(loss_fn)。
3. 定义gradient_descent_step函数。它应接受params（$x$的当前值）和learning_rate作为参数。
4. 在函数内部：
   • 使用grad_fn计算损失相对于params的梯度。
   • 计算更新后的参数：new_params = params - learning_rate * gradient_value。
5. 函数应返回new_params（这是更新后的状态）。
6. 初始化参数params（例如，jnp.array(5.0)）。
7. 设置learning_rate（例如，0.1）。
8. 在循环中执行几个更新步骤（例如，5步）：
   • 调用gradient_descent_step，传入当前的params和learning_rate。
   • 使用返回的new_params更新params变量。
   • 在每个步骤后打印参数值。
9. 可选： 对gradient_descent_step应用jax.jit。它能工作吗？为什么？

import jax
import jax.numpy as jnp

# 1. 定义损失函数
def loss_fn(x):
  return x**2

# 2. 获取梯度函数
grad_fn = jax.grad(loss_fn)

# 3. 定义状态更新函数（梯度下降步骤）
def gradient_descent_step(params, learning_rate):
  """执行一步梯度下降。"""
  gradient_value = grad_fn(params)
  new_params = params - learning_rate * gradient_value
  # 返回新状态（更新后的参数）
  return new_params

# 可选：JIT编译步骤函数
# jitted_gradient_descent_step = jax.jit(gradient_descent_step, static_argnums=(1,))
# 注意：如果learning_rate在编译函数的作用域内不改变，它通常被视为静态参数。

# 4. 初始化
current_params = jnp.array(5.0)
learning_rate = 0.1
num_steps = 5

print(f"初始参数: {current_params}")
print(f"学习率: {learning_rate}")

# 5. 执行更新步骤
for i in range(num_steps):
  # 传入当前状态 (参数)，获取新状态
  current_params = gradient_descent_step(current_params, learning_rate)
  # 如果正在JIT编译，请使用此行替代：
  # current_params = jitted_gradient_descent_step(current_params, learning_rate)
  print(f"步骤 {i+1}: 更新后的参数 = {current_params:.4f}")

预期输出：

初始参数: 5.0
学习率: 0.1
步骤 1: 更新后的参数 = 4.0000
步骤 2: 更新后的参数 = 3.2000
步骤 3: 更新后的参数 = 2.5600
步骤 4: 更新后的参数 = 2.0480
步骤 5: 更新后的参数 = 1.6384

本练习展示了状态（模型参数params）是如何在更新函数外部明确管理的。gradient_descent_step函数是纯粹的；它根据输入计算新状态并返回。这种模式在JAX中构建优化器或训练循环时非常重要。请注意，如果您jit此函数，如果learning_rate在编译函数的作用域内不改变，它通常最好标记为静态参数（使用static_argnums或static_argnames），因为这可以提升编译效率。

这些实际练习巩固了JAX所需的状态管理函数式方法。通过明确地传入和传出状态，您的函数保持纯粹并与JAX的强大变换（如jit、grad、vmap和pmap）兼容。这种模式可以从简单的计数器有效扩展到涉及神经网络参数和优化器统计信息的嵌套PyTrees等复杂状态。