1 files changed, 13 insertions, 0 deletions
diff --git a/source/slang/slang-check-type.cpp b/source/slang/slang-check-type.cpp
index d402dde03..6a8f802f7 100644
--- a/source/slang/slang-check-type.cpp
+++ b/source/slang/slang-check-type.cpp
@@ -320,6 +320,19 @@ namespace Slang
                 getSink()->diagnose(typeExp.exp, Diagnostics::cannotDefinePtrTypeToManagedResource);
             }
         }
+        
+        // Differentiable type checking.
+        // TODO: This can be super slow. Switch to caching the result asap.
+        if (this->m_parentFunc && 
+            this->m_parentFunc->findModifier<JVPDerivativeModifier>())
+        {
+            auto diffTypeContext = this->getShared()->innermostDiffTypeContext();
+            if (auto subtypeWitness = as<SubtypeWitness>(
+                tryGetInterfaceConformanceWitness(result, getASTBuilder()->getDifferentiableInterface())))
+            {
+                diffTypeContext->registerDifferentiableType((DeclRefType*)result, subtypeWitness);
+            }
+        }
 
         *outProperType = result;
         return true;