def row_sum

bendtorch-core@0.1.0.0 · src/tensor/elementwise.bend:74

import Base

74# The row sum PyTorch's log_softmax backward takes of a one-hot gradient row
75# whose target word is t. A row of one column holds t alone, and PyTorch's
76# sum of one word is that word, -0.0 kept; a wider row adds its +0.0 words,
77# so the sum is t + 0.0.
78def row_sum(single: Bool, t: F32) -> F32:
79  match single:
80    case True{}:
81      t
82    case False{}:
83      F32.add(t, 0.0)