neural_tangents.stax.ABRelu

neural_tangents.stax.ABRelu(a, b, do_stabilize=False)[source]

ABReLU nonlinearity, i.e. a * min(x, 0) + b * max(x, 0).

Parameters

  • a (float) – slope for x < 0.

  • b (float) – slope for x > 0.

  • do_stabilize (bool) – set to True for very deep networks.

Return type

Tuple[InitFn, ApplyFn, LayerKernelFn]

Returns

(init_fn, apply_fn, kernel_fn).