import torch


def compute_output_gradient_times_vec_intermediate(relu_network, x, z0):
    """
    Compute z and Wz defined in output_gradient_times_vector() function.
    Notice that when the relu network input is 0, we use only the right
    derivative.

    Return:
      z, Wz: Refer to output_gradient_times_vector()
      unique_flag: whether the output gradient is unique. Note the gradient
      is non-unique if an input to a ReLU unit is 0.
    """
    num_linear_layers = int((len(relu_network) + 1) / 2)
    z = [None] * num_linear_layers
    Wz = [None] * num_linear_layers
    z[0] = z0
    linear_layer_input = x
    unique_flag = True
    for linear_layer_count in range(num_linear_layers):
        Wz[linear_layer_count] = relu_network[
            2 * linear_layer_count].weight.data @ z[linear_layer_count]
        linear_layer_output = relu_network[2 * linear_layer_count](
            linear_layer_input)
        if linear_layer_count < num_linear_layers - 1:
            relu_layer = relu_network[2 * linear_layer_count + 1]
            linear_layer_input = relu_layer(linear_layer_output)
            beta = linear_layer_output >= 0
            c = 0 if isinstance(relu_layer,
                                torch.nn.ReLU) else relu_layer.negative_slope
            if torch.any(torch.abs(linear_layer_output) < 1E-10):
                unique_flag = False
            M = c + (1 - c) * beta
            z[linear_layer_count + 1] = M * Wz[linear_layer_count]
    return z, Wz, unique_flag