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)