text: class PReLU(tfb.Bijector):
      def __init__(self, alpha=0.5, validate_args=False, name="p_relu"):
          super(PReLU, self).__init__(
              forward_min_event_ndims=0,
              validate_args=validate_args,
              name=name)
          self.alpha = alpha
  
      def _forward(self, x):
          return tf.where(tf.greater_equal(x, 0), x, self.alpha * x)
  
      def _inverse(self, y):
          return tf.where(tf.greater_equal(y, 0), y, 1. / self.alpha * y)
  
      def _inverse_log_det_jacobian(self, y):
          I = tf.ones_like(y)
          J_inv = tf.where(tf.greater_equal(y, 0), I, 1.0 / self.alpha * I)
          log_abs_det_J_inv = tf.log(tf.abs(J_inv))
          return log_abs_det_J_inv