[Python] Softmax with polars Lazy Dataframe

I'm relatively new to using polars and it seems to be very verbose compared to pandas for what I would consider even relatively basic manipulations.

Case in point, the shortest way I could figure out doing a softmax over a lazy dataframe is the following:

import polars as pl

data = pl.DataFrame({'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]}).lazy()
cols = ['a','b','c']

data = data.with_columns([ pl.col(c).exp().alias(c) for c in cols]) # Exp all columns
data = data.with_columns(pl.sum_horizontal(cols).alias('sum')) # Get row sum of exps
data = data.with_columns([ (pl.col(c)/pl.col('sum')).alias(c) for c in cols ]).drop('sum')

data.collect()


Am I missing something and is there a shorter, more readable way of achieving this?

Continue reading...