import torch


def compute_jacobian(x, y, structured_tensor=False,
                     retain_graph=False):
    """Compute the Jacobian matrix of output with respect to input.

    If input and/or output have more than one dimension, the Jacobian of the
    flattened output with respect to the flattened input is returned if
    `structured_tensor` is `False`. If `structured_tensor` is `True`, the 
    Jacobian is structured in dimensions `[y_shape, flattened_x_shape]`.
    Note that `y_shape` can contain multiple dimensions.
    
    Args:
        x (list or torch.Tensor): Input tensor or sequence of tensors with the
            parameters to which the Jacobian should be computed. Important:
            the `requires_grad` attribute of input needs to be `True` while
            computing output in the forward pass.
        y (torch.Tensor): Output tensor with the values of which the
            Jacobian is computed.
        structured_tensor (bool): A flag indicating if the Jacobian should be
            structured in a tensor of shape `[y_shape, flattened_x_shape]`
            instead of `[flattened_y_shape, flattened_x_shape]`.

    Returns:
        (torch.Tensor): 2D tensor containing the Jacobian of output with
            respect to input if `structured_tensor` is `False`.
            If `structured_tensor` is `True`, the Jacobian is structured in a
            tensor of shape `[y_shape, flattened_x_shape]`.
    """
    if isinstance(x, torch.Tensor):
        x = [x]

    # Create the empty Jacobian.
    output_flat = y.view(-1)
    numel_input = 0
    for input_tensor in x:
        numel_input += input_tensor.numel()
    jacobian = torch.Tensor(y.numel(), numel_input)

    # Compute the Jacobian.
    for i, output_elem in enumerate(output_flat):
        if i == output_flat.numel() - 1:
            gradients = torch.autograd.grad(output_elem, x,
                                            retain_graph=retain_graph,
                                            create_graph=False,
                                            only_inputs=True)
        else:
            gradients = torch.autograd.grad(output_elem, x,
                                            retain_graph=True,
                                            create_graph=False,
                                            only_inputs=True)
        jacobian_row = torch.cat([g.view(-1).detach() for g in gradients])
        jacobian[i, :] = jacobian_row

    if structured_tensor:
        shape = list(y.shape)
        shape.append(-1) 
        jacobian = jacobian.view(shape)

    return jacobian