[AutoDiff] Fix substitution map remapping bug. (#25689)

Remap pullback reabstraction thunk substitution map during VJP generation.
Resolves TF-534.

Author: dan-zheng

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -3336,7 +3336,8 @@ class VJPEmitter final
           loweredPullbackType);
       auto *thunkRef = getBuilder().createFunctionRef(loc, thunk);
       pullback = getBuilder().createPartialApply(
-          ai->getLoc(), thunkRef, thunk->getForwardingSubstitutionMap(),
+          ai->getLoc(), thunkRef,
+          getOpSubstitutionMap(thunk->getForwardingSubstitutionMap()),
           {pullback}, actualPullbackType->getCalleeConvention());
     }
     pullbackValues[ai->getParent()].push_back(pullback);

test/AutoDiff/autodiff_diagnostics.swift

@@ -170,6 +170,23 @@ struct TF_305 : Differentiable {
   }
 }
 
+protocol TF_534_Layer : Differentiable {
+  associatedtype Input : Differentiable
+  associatedtype Output : Differentiable
+
+  @differentiable
+  func callAsFunction(_ input: Input) -> Output
+}
+struct TF_534_Tensor<Scalar> : Differentiable {}
+
+func TF_534<Model: TF_534_Layer>(
+  _ model: inout Model, inputs: Model.Input
+) -> TF_534_Tensor<Float> where Model.Output == TF_534_Tensor<Float> {
+  return valueWithPullback(at: model) { model -> Model.Output in
+    return model(inputs)
+  }.0
+}
+
 //===----------------------------------------------------------------------===//
 // Classes and existentials (not yet supported)
 //===----------------------------------------------------------------------===//