From 17c7163c2ae8fc290e70b43d8700b68ef18b1ee1 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Fri, 6 Oct 2023 14:03:18 -0700
Subject: Small type system fixes. (#3265)

---
 source/slang/slang-ir-check-differentiability.cpp | 5 +++++
 1 file changed, 5 insertions(+)

(limited to 'source/slang/slang-ir-check-differentiability.cpp')

diff --git a/source/slang/slang-ir-check-differentiability.cpp b/source/slang/slang-ir-check-differentiability.cpp
index b937fe052..ddb70d779 100644
--- a/source/slang/slang-ir-check-differentiability.cpp
+++ b/source/slang/slang-ir-check-differentiability.cpp
@@ -85,7 +85,12 @@ public:
         switch (func->getOp())
         {
         case kIROp_ForwardDifferentiate:
+            if (auto fwdDerivative = func->getOperand(0)->findDecoration<IRForwardDerivativeDecoration>())
+                return isDifferentiableFunc(fwdDerivative->getForwardDerivativeFunc(), level);
+            return isDifferentiableFunc(func->getOperand(0), level);
         case kIROp_BackwardDifferentiate:
+            if (auto bwdDerivative = func->getOperand(0)->findDecoration<IRUserDefinedBackwardDerivativeDecoration>())
+                return isDifferentiableFunc(bwdDerivative->getBackwardDerivativeFunc(), level);
             return isDifferentiableFunc(func->getOperand(0), level);
         default:
             break;
-- 
cgit v1.2.3