Use `@dr.wrap_ad` with a method of a class

[maxfrei750]

@Speierers Is it currently possible to use the @dr.wrap_ad decorator for a method of a class? A prominent example would be the forward method of a torch.nn.Module.

If it is currently not supported, is there an elegant workaround then?

Thanks.

[Speierers]

Hi @maxfrei750,

I am not sure to understand this request. Could you give a more concret example where you would use that decorator on a class method?

[maxfrei750]

Of course! Let's say I would like to input an image img that was rendered using mitsuba (i.e. a TensorXf) into a torch network net (typically a subclass torch.nn.Module), while keeping the gradients. Then I would call net.forward(img). The forward method expects a torch.Tensor, so it would be convenient to wrap it using @dr.wrap_ad. However, currently this results in an error, since the very first argument of forward is not of type TensorXf, but is the network itself (the self argument of forward).

[Speierers]

You should wrap the evaluation of your model instead:

dr.wrap_ad(source='drjit', target='torch')
def eval(img):
    return model(img)

[maxfrei750]

I see. I thought that should be avoided, because it violates the scope of the function. But as long as it works... :smile: Still, IMHO it would be neat if we could wrap methods directly.

[winterwindwang]

You should wrap the evaluation of your model instead:

dr.wrap_ad(source='drjit', target='torch')
def eval(img):
    return model(img)

When I use the dr.wrap() as a wrapper, I get an error "AttributeError: module 'drjit' has no attribute 'wrap_ad'". The drjit is installed automatically with pip install mitsuba. Moreover, I cannot run the tutorials here in the docs. Should I complie the drjit from scratch?

[maxfrei750]

IMHO, it would have been better to open a separate issue, since your problem is not related to the initial question.

When I use the dr.wrap() as a wrapper, I get an error "AttributeError: module 'drjit' has no attribute 'wrap_ad'".

This feature was only merged yesterday and there hasn't been a new release yet. Hence, it cannot be installed from pip yet.

Should I complie the drjit from scratch?

That's already close to the solution, but from my experience, you need to compile mitsuba from scratch, which includes the compilation of drjit. Compiling only drjit never worked for me.

[Speierers]

@maxfrei750 if you have trouble compiling Dr.Jit on its own could you maybe open an issue dedicated to that? It should work fine without Mitsuba.

I will close this it isn't an issue but rather a discussion. Please open a thread on the discussion page if you need more assistance on using dr.wrap_ad.

[maxfrei750]

I will close this it isn't an issue but rather a discussion.

True.

if you have trouble compiling Dr.Jit on its own could you maybe open an issue dedicated to that? It should work fine without Mitsuba.

Just to clarify: drjit compiled flawlessly, I just ran into problem when using it from mitsuba, when I didn't compile both. I opened an issue in the mitsuba repo a while ago (https://github.com/mitsuba-renderer/mitsuba3/issues/204) and @njroussel said that it often works best, to compile both together.