如何在Jax中从损失函数中返回一个值的字典？

import jax
import jax.numpy as jnp

def loss_fn(params, inputs, targets):
    # 模型的前向传播
    predictions = model(params, inputs)
    
    # 计算损失
    loss = jnp.mean(jnp.square(predictions - targets))
    
    # 返回一个值的字典
    return {'loss': loss}

# 使用损失函数计算损失
params = ...
inputs = ...
targets = ...
loss_dict = loss_fn(params, inputs, targets)

# 获取损失值
loss_value = loss_dict['loss']