RuntimeError: unwrapped_count > 0 INTERNAL ASSERT FAILED at "../aten/src/ATen/functorch/TensorWrapper.cpp":181, please report a bug to PyTorch. Should have at least one dead wrapper

[richardrl]

functorch fails when the function requires manipulating a scalar variable

def kkt(p1_tilde, p2_tilde, lambda_1_tilde, lambda_2_tilde, q_pos1, q_theta1):
        rot_mat = torch.Tensor([[torch.cos(q_theta), -torch.sin(q_theta)],
                                [torch.sin(q_theta), torch.cos(q_theta)]])
        rot_mat_inv = torch.linalg.inv(rot_mat)
        box1_A = A_1_canonical @ rot_mat_inv
        box1_b = b_1_canonical + A_1_canonical @ rot_mat_inv @ q_pos1
        D = 2
        K = torch.zeros(2*D + num_1_points + num_2_points)
        # lagrangian gradient wrt p1
        K[:D] = 2* (p1_tilde - p2_tilde) + box1_A.T @ lambda_1_tilde
        K[D:2*D] = -2* (p1_tilde - p2_tilde) + A_2_canonical.T @ lambda_2_tilde
        K[2*D:2*D+num_1_points] = lambda_1_tilde * (box1_A @ p1_tilde - box1_b)
        K[2*D+num_1_points:2*D+num_1_points+num_2_points] = lambda_2_tilde * (A_2_canonical @ p2_tilde - b_2_canonical)
        return K

The above error in the title comes from having the scalar variable theta

J_q_k_out = jacrev(lambda q_pos1, q_theta1: kkt(p1_tilde, p2_tilde, lambda_1_tilde, lambda_2_tilde, q_pos1, q_theta1), argnums=(0, 1))(q_pos, q_theta)

cc @ezyang @gchanan @zou3519 @Chillee @samdow @soumith @kshitij12345 @janeyx99

[maulberto3]

I got that error too in some other use case. Strange thing is that it was working fine, and then boom.

[zou3519]

This is unexpected. @maulberto3 do you have a script that demonstrates the problem?

@richardrl do you have sample inputs for us to try? Otherwise we can try random inputs until we hit the problem.

[zou3519]

Tentatively marking hi-pri due to assertion failure. First thing we should do is come up with a repro.

[kshitij12345]

Could be related to #95561. Will take a look.

[malfet]

Can you please run python -mtorch.utils.collect_env and post the results here?

[soulitzer]

High pri to get a repro