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.

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

Support higher axis numbers for dr.any and dr.all reductions #304

Draft
wants to merge 2 commits into
base: master
Choose a base branch
from

Conversation

zhestyatsky
Copy link

@zhestyatsky zhestyatsky commented Oct 22, 2024

Hello,

I have been experimenting with the latest drjit version and noticed that dr.all() and dr.any() operations seem to not be supported for Tensors when they are evaluated with axis argument greater than 0.

Reproducer below raises `RuntimeError: drjit.all(<drjit.cuda.TensorXb>): tensor type is not compatible with the requested reduction.

import drjit as dr
from drjit.cuda import TensorXb

v = TensorXb([[True, False, True],
              [True, True, True],
              [True, False, True],
              [True, True, True]
             ])
dr.all(v, axis=1)

Similar exception is raised when calling dr.any().

This PR aims to enable dr.all() and dr.any() operations for Tensors with axis parameter greater than 0.

if (axis_len == 1 && red_axis == 0)
return reduce_seq(op, h, nb::int_(0), mode);
nb::raise_type_error("tensor type is not compatible with "
"the requested reduction.");
}
// Complex case, defer to a separate Python implementation
ReduceOp reduceOp;
Copy link
Member

Choose a reason for hiding this comment

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

Very minor. rendceOp -> reduce_op (underscore case for vriables)

@wjakob
Copy link
Member

wjakob commented Jan 28, 2025

Hi @zhestyatsky -- is this PR ready to be merged or still under development? It states "Draft" status.
One thing I am surprised about is that this works as-is.

The all/and reductions are AFAIK only implemented for 32-bit integers. We previously got away with using them for reduction down to 1 element by applying a simple follow-up 4->1 element reduction and doing something some extra memsets in case the input array wasn't exactly a multiple of 4 elements.

For reductions along another axes, I am wondering if there won't be mix-ups when the element count isn't a multiple of 4?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants