code: def ParametricRelu(a=1.0):
      r"""Returns a layer that computes a ReLU function with the given slope.
  
    .. math::
        f(x) = \left\{ \begin{array}{cl}
            0  & \text{if}\ x \leq 0, \\
            ax & \text{otherwise}.
        \end{array} \right.
  
    Args:
      a: Slope of line for positive inputs.
    """
      return Fn(
          "ParametricRelu",
          lambda x: jnp.maximum(a * x, jnp.zeros_like(x)),
      )
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
label_desc: Use image node_id 2 for calling a global function with example usage: ParametricRelu(a) and returns: Fn