element_wise

class braintrace.element_wise(weight, fn=<function <lambda>>)

ETP-aware element-wise operation.

Applies fn to weight and passes through a marker primitive. The operation is treated as diagonal in the hidden-state space.

Parameters:

• weight – Weight parameter.

• fn – Element-wise function. Defaults to identity.

Returns:

fn(weight).