Simplifying Softmax Calculation for a Polars Lazy DataFrame

Polars’ lazy API is designed to let you build up an explicit computation plan that later gets optimized and executed. That’s why many basic manipulations may seem more verbose compared to the in‑place operations of pandas. In your softmax example you intentionally “bake in” the transformations as separate steps, which gives you clear control but at the expense of brevity.

There isn’t a built‑in softmax function in Polars, so you need to express the steps (exponentiation, row‑sum, and normalization) explicitly. However, you can refactor your code a little to make it a single chained pipeline or encapsulate the logic as a helper function.

Single Pipeline Example

You can combine all the steps into one chain:

PYTHONCopied
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']

result = (
    data.with_columns([pl.col(c).exp().alias(c) for c in cols])         # Compute exp(c)
        .with_columns(pl.sum_horizontal(cols).alias("sum"))               # Row-wise sum of exponentials
        .with_columns([(pl.col(c) / pl.col("sum")).alias(c) for c in cols]) # Normalize columns
        .drop("sum")
        .collect()
)

print(result)

Helper Function Example

If you find yourself writing this pattern repeatedly, you might wrap it in a helper function:

PYTHONCopied
def softmax(df: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
    return (
        df.with_columns([pl.col(c).exp().alias(c) for c in columns])
          .with_columns(pl.sum_horizontal(columns).alias("sum"))
          .with_columns([(pl.col(c) / pl.col("sum")).alias(c) for c in columns])
          .drop("sum")
    )

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()

result = softmax(data, ['a', 'b', 'c']).collect()
print(result)

Summary

• Verbosity: It might seem verbose compared to pandas, but that’s a conscious design choice so that every transformation is explicit and optimal.
• Readability: Chaining the operations or putting them in a small helper function can help make your code more readable.
• Performance: The verbosity also offers clarity when building optimized lazy query plans, which is a major strength of polars.

So, while there isn’t a “shorter” way per se—because the operations you’re performing are genuinely multi-step—you have some flexibility in how you organize the code without sacrificing clarity or performance.