code: def _w_delta_squared(z, delta):
      """Applies W_delta transformation to the input.
  
      For a given z, `W_delta(z) = sign(z) * (W(delta * z^2)/delta)^0.5`. This
      transformation is defined in Equation (9) of [1].
  
      Args:
        z: Input of the transformation.
        delta: Parameter delta of the transformation.
  
      Returns:
        The transformed Tensor with same shape and same dtype as `z`.
      """
      delta = tf.convert_to_tensor(delta, dtype=z.dtype)
      z = tf.broadcast_to(
          z, ps.broadcast_shape(ps.shape(z), ps.shape(delta))
      )
      wd = tf.sign(z) * tf.sqrt(
          special.lambertw(delta * z**2) / delta
      )
      return tf.where(tf.equal(delta, 0.0), z, wd)
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.math import special
label_desc: Use image node_id 2 for calling a global function with example usage: _w_delta_squared(z, delta) and returns: tf