docs: def _nested_unary_mul(nested_a, p):
    """Multiply `Tensors` in arbitrarily nested `Tensor` `nested_a` with `p`."""
    def mul_with_broadcast(tensor):
      ndims = tensor.shape.ndims
      if ndims != 2:
        p_reshaped = tf.reshape(p, [-1] + [1] * (ndims - 1))
        return p_reshaped * tensor
      else:
        return p * tensor
    return nest.map(mul_with_broadcast, nested_a)