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).