Jax加权交叉熵损失广播形状不兼容问题求助

问题重现 #

尝试用Jax编写加权交叉熵损失函数时遇到维度不兼容报错，代码如下：

import jax.numpy as np
from functools import partial
import jax

@partial(np.vectorize, signature="(c),(),()->()")
def weighted_cross_entropy_loss(logits, label, weights):
    one_hot_label = jax.nn.one_hot(label, num_classes=logits.shape[0])
    return -np.sum(weights* logits*one_hot_label)

logits=np.array([[1,2,3,4,5,6,7],[2,3,4,5,6,7,8]])
labels=np.array([1,2])
weights=np.array([1,2,3,4,5,6,7])
print(weighted_cross_entropy_loss(logits,labels,weights))

报错信息：

Traceback (most recent call last):
  File "/opt/homebrew/anaconda3/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 147, in broadcast_shapes
    return _broadcast_shapes_cached(*shapes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/anaconda3/lib/python3.11/site-packages/jax/_src/util.py", line 284, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/anaconda3/lib/python3.11/site-packages/jax/_src/util.py", line 277, in cached
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/anaconda3/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 153, in _broadcast_shapes_cached
    return _broadcast_shapes_uncached(*shapes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/anaconda3/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 169, in _broadcast_shapes_uncached
    raise ValueError(f"Incompatible shapes for broadcasting: shapes={list(shapes)}")
ValueError: Incompatible shapes for broadcasting: shapes=[(2,), (2,), (7,)]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/PATH/test.py", line 15, in <module>
    print(weighted_cross_entropy_loss(logits,labels,weights))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/anaconda3/lib/python3.11/site-packages/jax/_src/numpy/vectorize.py", line 274, in wrapped
    broadcast_shape, dim_sizes = _parse_input_dimensions(
                                 ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/anaconda3/lib/python3.11/site-packages/jax/_src/numpy/vectorize.py", line 123, in _parse_input_dimensions
    broadcast_shape = lax.broadcast_shapes(*shapes)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/anaconda3/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 149, in broadcast_shapes
    return _broadcast_shapes_uncached(*shapes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/anaconda3/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 169, in _broadcast_shapes_uncached
    raise ValueError(f"Incompatible shapes for broadcasting: shapes={list(shapes)}")
ValueError: Incompatible shapes for broadcasting: shapes=[(2,), (2,), (7,)]

问题原因 #

• np.vectorize签名不匹配：设置的签名"(c),(),()->()"要求输入为(c,)维度数组、标量、标量，但实际传入的weights是(7,)维度数组，Jax广播时无法匹配维度，触发报错。
• 不必要的vectorize使用：Jax数组操作本身支持向量化，手动包装反而会引入维度处理冗余问题。
• 交叉熵计算逻辑错误：直接用原始logits计算不符合交叉熵定义，需要先通过log_softmax转换为对数概率。

修正后的代码 #

import jax.numpy as np
import jax

def weighted_cross_entropy_loss(logits, labels, weights):
    # 对logits计算对数概率，保持(2,7)维度
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    # 将标签转为one-hot编码，维度变为(2,7)
    one_hot_labels = jax.nn.one_hot(labels, num_classes=logits.shape[-1])
    # 计算每个样本的加权损失，用点积替代sum提升效率
    sample_losses = -np.sum(weights * log_probs * one_hot_labels, axis=-1)
    # 返回所有样本损失的总和（可根据需求改为均值）
    return np.sum(sample_losses)

logits = np.array([[1,2,3,4,5,6,7],[2,3,4,5,6,7,8]])
labels = np.array([1,2])
weights = np.array([1,2,3,4,5,6,7])
print(weighted_cross_entropy_loss(logits, labels, weights))

关键说明 #

• 移除np.vectorize后，函数直接处理批量输入，维度匹配更自然；
• jax.nn.log_softmax确保计算符合交叉熵的数学定义；
• weights为(7,)维度时，Jax会自动广播到(2,7)维度，完成逐类别加权；
• 最终返回的损失总和可根据训练需求调整为均值。

内容的提问来源于stack exchange，提问作者WillWu