Backward not implemented for tinygrad.mlops.Eq

[davej23]

Hi everyone,

I tried to find previous posts (issues + discussions) about this but couldn't find anything - if I have missed a previous discussion about this feel free to just point me to it 😄

I am looking at implementing some gradient-based explainability methods, such as GradCAM and Integrated Gradients, for Tinygrad.

For explaining object detection models you need a wrapper function which takes the output (shape (1, 4 + num classes, 86400) for YOLOv8 and converts it into something like (1, N, 5 + num classes) where the extra value is the 'objectness' or highest probability in the classification probability vector. After this, we can onehot encode the class of the ground truth (1, M, 5) to get (1, M, 5 + num classes) and then compare this to the output of the wrapper to get a single scalar 'similarity' score of the prediction to the ground truth.

I have implemented this and when I run .backward() on the score function, using something as simple as (wrapper_output - ground_truth).sum(), I get

**RuntimeError: backward not implemented for <class 'tinygrad.mlops.Eq'>**

error in the wrapper function below due to the top_idxs = x[:, 4:, :].max(1).argmax(1) line. If I create a new tensor instead, the line below it, it works and I don't get the error.

This is the code I had set up for the wrapper function

def wrapper(x: Tensor) -> Tensor: 
    """This function selects top box by 'objectness' (box with highest probability) """
    top_idxs = x[:, 4:, :].max(1).argmax(1)  # This doesn't work, 'Backward not implemented error'
    # top_idxs = Tensor(x[:, 4:, :].max(1).argmax(1).numpy())  # This works, no 'Backward not implemented error'
    top_x = Tensor.stack([x[i, :, top_idxs[i]].unsqueeze(1) for i in range(x.shape[0])])  # (BS, 4 + NC, 1)
    x_fast = top_x[:, :4, :]
    x_fast = x_fast.cat(top_x[:, 4:, :].max(1).unsqueeze(-1), dim=1)
    x_fast = x_fast.cat(top_x[:, 4:, :], dim=1)
    x_fast = x_fast.permute((0, 2, 1))
    return x_fast

I was just wanting to understand why this error occurs, thanks! 😄

[chenyuxyz]

The error occurs because argmax is not differentiable and has no gradient.

[davej23]

Ah of course, that makes sense, thanks for the quick response!