Differentiable generic functions and type constraints

[MangoSister]

Hello, I have just started using slang for its autodiff capability so this might be trivial or a non-issue. This is with slang 2024.11 and Vulkan, but I don't see any related change since 2024.11. I am confused with the behavior of differentiable generic functions. In my project, I've defined a bunch of generic math utilities functions such as:

T sqr<T>(T x)
    where T : __BuiltinArithmeticType { return x * x; }

Then, I wanted to write a differentiable function that called sqr so I marked it as differentiable and wrote something like this (simplified):

[Differentiable]
T sqr<T>(T x)
    where T : __BuiltinArithmeticType { return x * x; }

//...

[Differentiable]
float foo(float x, no_diff float a)
{
    float d = x - a;
    return sqr(d);
}

//...
var xd = diffPair(x); 
bwd_diff(foo)(xd, a, 1.0);
// do something with xd.d...

However, I found that the derivative xd.d was always zero. This seemed innocuous until I realized that maybe sqr was not working correctly because __BuiltinArithmeticType includes integer types that are not differentiable. So I created the following explicit overloaded versions for both integer and float types:

// Overload for float types.
[Differentiable]
T sqr<T>(T x)
    where T : __BuiltinFloatingPointType { return x * x; }

// Overload for integer types that is not marked as differentiable.
T sqr<T>(T x)
    where T : __BuiltinIntegerType { return x * x; }

And then foo could be autodiff-ed correctly.

My question is: Is this the intended behavior? It seems to me that when marking generic functions that accept both differentiable and non-differentiable types as differentiable, one should expect one of the following two reasonable outcomes:

1. The function produces correct derivatives for differentiable types, and does nothing extra for non-differentiable types.
2. This should be explicitly disallowed. Only generic functions that accept IDifferentiable are allowed to be marked as differentiable. So users are forced to write something like my above example.

However, currently neither of the above happens and we simply get a silent failure even with differentiable types (float). This seems to be confusing (and took me a long while to figure out...) Thanks.

[saipraveenb25]

Thank you for opening a discussion on this! I understand that having gradients disappear silently can be very annoying to track down. We do have a gradient data-flow analysis pass to warn about these cases (i.e. requires the 'no_diff' modifier when calling a non-differentiable function with differentiable inputs). It looks like this case doesn't get detected because the function is marked [Differentiable]

This is a slightly tricky situation. You are correct in that its confusing that

[Differentiable]
T sqr<T>(T x)
    where T : __BuiltinArithmeticType { return x * x; }

is allowed but silently returns a 0 since __BuiltinArithmeticType is non-differentiable.

One fix is to warn users if the function has no side-effects & no differentiable inputs/outputs.

However, more complicated cases with a mix of differentiable and non-differentiable types will still pose a problem

[Differentiable]
T foo<T, U>(T x, U y)
    where T : __BuiltinArithmeticType 
    where U : __BuiltinFloatingPointType { /* .. */ }

Call foo(a, b) where you expect a to have a derivative will still cause a silent failure since there won't be a warning.

A more solid fix would be to extend the no_diff system to error out when passing a differentiable type to a non-differentiable parameter, and require that the value be wrapped in a no_diff qualifier
Right now, the no_diff system only tests the differentiability of the function itself, but not the parameter being passed to.

[saipraveenb25]

Opened an issue so we can track the fix: #5754

[MangoSister]

Thank you so much Sai for the explanation. Yes, I agree that extending no_diff should be the way to go. Hope there could be a fix for this soon!