BF16 squared ReLU and its backward are ordinary compositions. Upstream
rounds both ReLU's result and the square; backward rounds 2*relu(x), then
its product with dy, then applies threshold-backward. Multiplying by a
zero mask would propagate NaNs and is not the same derivative program.
367def pointwiseReluSquared =
368 (lambda erased inputs : Nat .
369 (lambda unrestricted x : (family PointwiseExpression StdU32 inputs) .
370 (let unrestricted relu =
371 (pointwiseUnary
372 StdU32
373 inputs
374 (constructor PointwiseUnary PointwiseRoundBF16)
375 (pointwiseBinary StdU32 inputs (constructor PointwiseBinary PointwisePositiveGate) x x))
376 in
377 (pointwiseUnary
378 StdU32
379 inputs
380 (constructor PointwiseUnary PointwiseRoundBF16)
381 (pointwiseUnary StdU32 inputs (constructor PointwiseUnary PointwiseSquare) relu)))))The compiler supplied declaration spans and resolved links from this source snapshot. This page does not assert that this file belongs to a checked closure.