Invalid HLSL generated for bwd_diff function

[tunabrain]

I'm running into an issue where slang will generate invalid HLSL for the following snippet:

struct Foo
{
    float foo(int idx) { return 5.0f; }
}
struct WrappedBuffer<T, let D : int>
{
    StructuredBuffer<T> buffer;
    int[D] shape;

    T get(int idx) { return buffer[idx]; }
}
struct GradInBuffer<T : IDifferentiable, let D : int>
{
    RWStructuredBuffer<T> primal;
    WrappedBuffer<T.Differential, D> grad_in;
}

[Differentiable, BackwardDerivative(set_bwd)]
void set(GradInBuffer<float[1], 1> t, int idx, float[1] value)
{
    t.primal[idx] = detach(value);
}
void set_bwd(GradInBuffer<float[1], 1> t, int idx, inout DifferentialPair<float[1]> grad)
{
    // Generates invalid HLSL, "candidate function not viable: no known conversion from 'WrappedBuffer_0' to 'WrappedBuffer_1' for 1st argument"
    grad = diffPair(grad.p, t.grad_in.get(idx));

    // This works:
    //grad = diffPair(grad.p, t.grad_in.buffer[idx]);
}

struct CallData
{
    Foo weights;
    RWStructuredBuffer<float> dOut;
    GradInBuffer<float[1], 1> result;
}
ParameterBlock<CallData> call_data;

[shader("compute")]
[numthreads(32, 1, 1)]
void main(uint3 dispatchThreadID: SV_DispatchThreadID)
{
    int idx = dispatchThreadID.x;

    float[1] primal = { call_data.weights.foo(idx) };
    DifferentialPair<float[1]> dResult = diffPair(primal);
    bwd_diff(set)(call_data.result, idx, dResult);

    call_data.dOut[idx] = dResult.d[0];
}

It seems to generate two identical specializations for WrappedBuffer (WrappedBuffer_0 and WrappedBuffer_1), and calls the get method of WrappedBuffer_1 with an instance of WrappedBuffer_0. This only happens for the backwards pass and is very finicky; e.g. removing int[D] shape; or changing T from float[1] to float makes it go away.