Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fixing out of bound access for nll_loss #1752

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
Open

fixing out of bound access for nll_loss #1752

Changes from 2 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
01e0c4e
wip
jjsjann123 Feb 7, 2025
3edce5b
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 7, 2025
178ca95
wip
jjsjann123 Feb 7, 2025
7c78ab7
Merge remote-tracking branch 'origin/cross_entropy_out_of_bound' into…
jjsjann123 Feb 7, 2025
e9d6e40
Merge branch 'main' into cross_entropy_out_of_bound
jjsjann123 Feb 7, 2025
fde1b3f
WIP
jjsjann123 Feb 7, 2025
f0e814d
fixing handling of weight
jjsjann123 Feb 7, 2025
2eb2c77
fixing test
jjsjann123 Feb 8, 2025
419dd05
adding comment
jjsjann123 Feb 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions thunder/torch/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5111,12 +5111,18 @@ def _nll_loss_helper(
bcast_weight = reshape(weight, [num_class] + [1 for _ in range(2, a.ndim)])
out = out * bcast_weight

assert isinstance(ignore_index, Number)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If ignore_index is a NumberProxy are the constraints created as expected for >= 0 and < num_class comparisons with the symbolic values cache mode?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For today with constant constraints, I think that's right. But then we could have the same out of bound access issue with NumberProxy for ignore_index.
If we do need to support that, I think for now we need to use < 0 or >= num_class to handle it.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

BTW, added the support.

For some reason thunderfx does indeed hit give ignore_index as NumberProxy instead, you probably know this better than I do.


# Make target broadcastable with output, which has same shape as input tensor.
bcast_target = unsqueeze(target, class_dim)

out = take_along_dim(out, bcast_target, class_dim)
selected_target_mask = bcast_target != ignore_index
out = where(selected_target_mask, out, 0)
if ignore_index >= 0 and ignore_index < num_class:
out = take_along_dim(out, bcast_target, class_dim)
selected_target_mask = bcast_target != ignore_index
out = where(selected_target_mask, out, 0)
else:
selected_target_mask = bcast_target != ignore_index
index = where(selected_target_mask, bcast_target, num_class)
padded_out = clang.pad(out, 0, [1] * (class_dim - 1) * 2 + [0, 1] + [1] * (out.ndim - class_dim) * 2)

# This section handles applying the reduction parameter to the output.
# We return None for the total_weight when reduction is "none" or "sum" since it is unused in the backwards pass.
Expand Down
Loading