如何通过Polars表达式实现LazyFrame的Softmax运算？是否存在更简洁的写法？

刚接触Polars的时候，确实会觉得它的写法和Pandas比起来有点啰嗦，尤其是这种需要多步操作的计算～不过其实咱们可以利用Polars表达式的链式组合能力，把你原来的三步操作压缩成一步，代码会简洁很多，可读性也更好。

简化后的完整实现 #

import polars as pl

df = pl.LazyFrame({'a': [1,2,3,4,5,6,7,8,9,10], 'b':[5,5,5,5,5,5,5,5,5,5], 'c': [10,9,8,7,6,5,4,3,2,1]})
cols = ['a','b','c']

# 一步完成Softmax计算
df = df.with_columns(
    (pl.col(cols).exp() / pl.col(cols).exp().sum_horizontal()).alias(cols)
)

print(df.collect())

代码逻辑说明 #

这个写法把“计算指数→求行和→做除法”三个步骤合并成了一个表达式：

• pl.col(cols).exp()：对指定列计算指数，和你第一步的操作完全一致；
• .sum_horizontal()：直接对指数化后的列按行求和，替代原来单独创建sum列的步骤；
• 最后用指数化的列除以行和，直接得到Softmax结果，再通过.alias(cols)把结果命名回原来的列名，省掉了后续的drop('sum')操作。

运行结果（和原写法完全一致） #

shape: (10, 3)
┌──────────┬──────────┬──────────┐
│ a        ┆ b        ┆ c        │
│ ---      ┆ ---      ┆ ---      │
│ f64      ┆ f64      ┆ f64      │
╞══════════╪══════════╪══════════╡
│ 0.000123 ┆ 0.006692 ┆ 0.993185 │
│ 0.000895 ┆ 0.01797  ┆ 0.981135 │
│ 0.006377 ┆ 0.047123 ┆ 0.946499 │
│ 0.04201  ┆ 0.114195 ┆ 0.843795 │
│ 0.211942 ┆ 0.211942 ┆ 0.576117 │
│ 0.576117 ┆ 0.211942 ┆ 0.211942 │
│ 0.843795 ┆ 0.114195 ┆ 0.04201  │
│ 0.946499 ┆ 0.047123 ┆ 0.006377 │
│ 0.981135 ┆ 0.01797  ┆ 0.000895 │
│ 0.993185 ┆ 0.006692 ┆ 0.000123 │
└──────────┴──────────┴──────────┘

进阶：封装成可复用函数 #

如果需要经常用到Softmax运算，还可以把逻辑封装成自定义函数，调用起来更顺手：

def softmax(cols):
    exp_vals = pl.col(cols).exp()
    return exp_vals / exp_vals.sum_horizontal()

# 调用方式更简洁
df = df.with_columns(softmax(cols).alias(cols))

你原来的写法完全是正确的，只是没有充分利用Polars表达式的嵌套和链式调用特性——这种特性可以让多步操作合并成一个连贯的表达式，既减少代码行数，也能让逻辑更清晰。

备注：内容来源于stack exchange，提问作者velochy