Would be useful to have an option for specifying a custom epilog a la custom inits. Wrinkles:

• pointwise would be simple enough - x => ${expr} - but "real" epilogs (layernorm, scaled softmax) are plumbed in kind of sloppily currently, would need a revamp
• param logistics, there's already a top level property called expr. epilog_expr probably