Issue Description:

We had been maintaining a temporary branch during the holiday commit freeze during builds of the Swift for TensorFlow toolchain, and after switching back to main after the holidays we're now experiencing an assertion failure during the Differentiation SILModuleTransform for both the tests on tensorflow/swift-apis and the models within tensorflow/swift-models.

The assertion is the same in both cases, and produces the following failure:

swift-frontend: /swift-base/swift/lib/SILOptimizer/Utils/DifferentiationMangler.cpp:50: swift::Demangle::NodePointer mangleAsNode(swift::SILFunction *, Demangle::AutoDiffFunctionKind, swift::AutoDiffConfig, swift::Demangle::Demangler &): Assertion `demangledOrig->getNumChildren() == 1' failed.
Please submit a bug report (https://swift.org/contributing/#reporting-bugs) and include the project and the crash backtrace.
Stack dump:
0. Program arguments: /swift-tensorflow-toolchain/usr/bin/swift-frontend -frontend -c /swift-apis/Tests/TensorFlowTests/CollectionTesting.swift /swift-apis/Tests/TensorFlowTests/ContextTests.swift /swift-apis/Tests/TensorFlowTests/EpochsTests.swift /swift-apis/Tests/TensorFlowTests/FreezableTests.swift /swift-apis/Tests/TensorFlowTests/Helpers.swift /swift-apis/Tests/TensorFlowTests/InitializerTests.swift /swift-apis/Tests/TensorFlowTests/KeyPathIterableTests.swift -primary-file /swift-apis/Tests/TensorFlowTests/LayerTests.swift /swift-apis/Tests/TensorFlowTests/LazyTensorEvaluationTests.swift /swift-apis/Tests/TensorFlowTests/LazyTensorExplicitTraceTests.swift /swift-apis/Tests/TensorFlowTests/LazyTensorHandleTests.swift /swift-apis/Tests/TensorFlowTests/LazyTensorOperationTests.swift /swift-apis/Tests/TensorFlowTests/LazyTensorShapeInferenceTests.swift /swift-apis/Tests/TensorFlowTests/LazyTensorTestHelper.swift /swift-apis/Tests/TensorFlowTests/LazyTensorTFFunctionBuilderTests.swift /swift-apis/Tests/TensorFlowTests/LazyTensorTraceCacheTests.swift /swift-apis/Tests/TensorFlowTests/LazyTensorTraceTests.swift /swift-apis/Tests/TensorFlowTests/LossTests.swift /swift-apis/Tests/TensorFlowTests/OptimizerTests.swift /swift-apis/Tests/TensorFlowTests/RuntimeTests.swift /swift-apis/Tests/TensorFlowTests/SamplingTests.swift /swift-apis/Tests/TensorFlowTests/SequencedTests.swift /swift-apis/Tests/TensorFlowTests/SequentialTests.swift /swift-apis/Tests/TensorFlowTests/SlicesTests.swift /swift-apis/Tests/TensorFlowTests/TensorAutoDiffTests.swift /swift-apis/Tests/TensorFlowTests/TensorGroupTests.swift /swift-apis/Tests/TensorFlowTests/TensorTests.swift /swift-apis/Tests/TensorFlowTests/TrivialModelTests.swift /swift-apis/Tests/TensorFlowTests/XCTestManifests.swift /swift-apis/Tests/TensorFlowTests/OperatorTests/BasicTests.swift /swift-apis/Tests/TensorFlowTests/OperatorTests/ComparisonTests.swift /swift-apis/Tests/TensorFlowTests/OperatorTests/ImageTests.swift /swift-apis/Tests/TensorFlowTests/OperatorTests/LinearAlgebraTests.swift /swift-apis/Tests/TensorFlowTests/OperatorTests/MathTests.swift /swift-apis/Tests/TensorFlowTests/OperatorTests/MatrixTests.swift /swift-apis/Tests/TensorFlowTests/OperatorTests/NNTests.swift -emit-module-path Tests/TensorFlowTests/CMakeFiles/TensorFlowTests.dir/LayerTests.swift.swiftmodule -emit-module-doc-path Tests/TensorFlowTests/CMakeFiles/TensorFlowTests.dir/LayerTests.swift.swiftdoc -emit-module-source-info-path Tests/TensorFlowTests/CMakeFiles/TensorFlowTests.dir/LayerTests.swift.swiftsourceinfo -emit-dependencies-path Tests/TensorFlowTests/CMakeFiles/TensorFlowTests.dir/LayerTests.swift.o.d -emit-reference-dependencies-path Tests/TensorFlowTests/CMakeFiles/TensorFlowTests.dir/LayerTests.swift.o.swiftdeps -target x86_64-unknown-linux-gnu -disable-objc-interop -I /swift-apis/Sources/CTensorFlow/include -I swift -I libtensorflow-prefix/src/libtensorflow -I libtensorflow-prefix/src/libtensorflow/tensorflow -I libtensorflow-prefix/src/libtensorflow/bazel-bin -I libtensorflow-prefix/src/libtensorflow/bazel-libtensorflow/external/com_google_absl -I libtensorflow-prefix/src/libtensorflow/bazel-libtensorflow/external/com_google_protobuf/src -I libtensorflow-prefix/src/libtensorflow/bazel-libtensorflow/external/eigen_archive -I swift-numerics-prefix/src/swift-numerics-build/swift -I swift-numerics-prefix/src/swift-numerics/Sources/_NumericsShims/include -module-link-name TensorFlowTests -O -D TENSORFLOW_USE_STANDARD_TOOLCHAIN -D TensorFlowTests_EXPORTS -Xllvm -sil-inline-generics -Xllvm -sil-partial-specialization -parse-as-library -module-name TensorFlowTests -o Tests/TensorFlowTests/CMakeFiles/TensorFlowTests.dir/LayerTests.swift.o

1. Swift version 5.4-dev (LLVM 3bb55ca, Swift f205b20)
2. While evaluating request ExecuteSILPipelineRequest(Run pipelines { Mandatory Diagnostic Passes + Enabling Optimization Passes } on SIL for TensorFlowTests.TensorFlowTests)
3. While running pass #​9129 SILModuleTransform "Differentiation".
4. While canonicalizing differentiable_function SIL node %611 = differentiable_function [parameters 0 1 2] [results 0] %610 : $@callee_guaranteed (@guaranteed Array<Tensor>, @guaranteed LSTMCell.State, @guaranteed LSTMCell.State) -> @owned Tensor // users: %632, %612
5. While ...in SIL function "@$s15TensorFlowTests05LayerC0C31testBidirectionalLSTMStackMergeyyFyyXEfU_".
   for expression at [/swift-apis/Tests/TensorFlowTests/LayerTests.swift:4130:28 - line:4225:5] RangeText="{
   var lstm = BidirectionalLSTM(
   LSTMCell(inputSize: 3, hiddenSize: 4),
   mergeFunction: LSTMCell.State.stack
   )

lstm.forward.cell.fusedWeight =
swapForgetUpdate(Tensor(concatenating: [kernelForward, recurrentKernelForward]))
lstm.forward.cell.fusedBias = swapForgetUpdate(biasForward)

lstm.backward.cell.fusedWeight =
swapForgetUpdate(Tensor(concatenating: [kernelBackward, recurrentKernelBackward]))
lstm.backward.cell.fusedBias = swapForgetUpdate(biasBackward)

let inputs = x.squeezingShape(at: 0).unstacked().map { $0.rankLifted() }

let initialForwardLayerState = LSTMCell.State(
cell: initialStateCellForward,
hidden: initialStateHiddenForward)
let initialBackwardLayerState = LSTMCell.State(
cell: initialStateCellBackward,
hidden: initialStateHiddenBackward)

let outputs = lstm(
inputs,
initialForwardLayerState: initialForwardLayerState,
initialBackwardLayerState: initialBackwardLayerState
).map { $0.hidden.unstacked(alongAxis: 0) }
assertEqual(
Tensor(concatenating: outputs.map { $0[0] }),
expectedStatesForward,
accuracy: 1e-6)
assertEqual(
Tensor(concatenating: outputs.map { $0[1] }),
expectedStatesBackward,
accuracy: 1e-6)

let (outputSum, (gradLSTM)) =
valueWithGradient(at: lstm) { lstm in
lstm.lastOutput(
from: inputs,
initialForwardLayerState: initialForwardLayerState,
initialBackwardLayerState: initialBackwardLayerState
).hidden.sum()
}
assertEqual(
outputSum,
expectedSum,
accuracy: 1e-6)
assertEqual(
gradLSTM.forward.cell.fusedWeight,
swapForgetUpdate(Tensor(concatenating: [expectedGradKernelForward, expectedGradRecurrentKernelForward])),
accuracy: 1e-6)
assertEqual(
gradLSTM.forward.cell.fusedBias,
swapForgetUpdate(expectedGradBiasForward),
accuracy: 1e-6)
assertEqual(
gradLSTM.backward.cell.fusedWeight,
swapForgetUpdate(Tensor(concatenating: [expectedGradKernelBackward, expectedGradRecurrentKernelBackward])),
accuracy: 1e-6)
assertEqual(
gradLSTM.backward.cell.fusedBias,
swapForgetUpdate(expectedGradBiasBackward),
accuracy: 1e-6)

let (gradInputs, gradInitialStateForward, gradInitialStateBackward) =
gradient(at: inputs, initialForwardLayerState, initialBackwardLayerState) {
inputs, initialForwardLayerState, initialBackwardLayerState in
lstm.lastOutput(
from: inputs,
initialForwardLayerState: initialForwardLayerState,
initialBackwardLayerState: initialBackwardLayerState
).hidden.sum()
}
assertEqual(
Tensor(concatenating: gradInputs.map { $0 }),
expectedGradX,
accuracy: 1e-6)
assertEqual(
gradInitialStateForward.hidden,
expectedGradInitialStateHiddenForward,
accuracy: 1e-6)
assertEqual(
gradInitialStateForward.cell,
expectedGradInitialStateCellForward,
accuracy: 1e-6)
assertEqual(
gradInitialStateBackward.hidden,
expectedGradInitialStateHiddenBackward,
accuracy: 1e-6)
assertEqual(
gradInitialStateBackward.cell,
expectedGradInitialStateCellBackward,
accuracy: 1e-6)
"
6. While processing // differentiability witness for specialized closure #​6 in closure #​1 in LayerTests.testBidirectionalLSTMStackMerge()
sil_differentiability_witness private [parameters 0 1 2] [results 0] @$s15TensorFlowTests05LayerC0C31testBidirectionalLSTMStackMergeyyFyyXEfU_0aB00A0VySfGSayAHG_AE8LSTMCellV5StateVySf_GANtcfU4_Tf2nnni_n : $@convention(thin) (@guaranteed Array<Tensor>, @guaranteed LSTMCell.State, @guaranteed LSTMCell.State, @guaranteed BidirectionalRecurrentLayer<LSTMCell>) -> @owned Tensor {
}

on SIL function "@$s15TensorFlowTests05LayerC0C31testBidirectionalLSTMStackMergeyyFyyXEfU_0aB00A0VySfGSayAHG_AE8LSTMCellV5StateVySf_GANtcfU4_Tf2nnni_n".
for expression at [/swift-apis/Tests/TensorFlowTests/LayerTests.swift:4197:83 - line:4204:9] RangeText="{
inputs, initialForwardLayerState, initialBackwardLayerState in
lstm.lastOutput(
from: inputs,
initialForwardLayerState: initialForwardLayerState,
initialBackwardLayerState: initialBackwardLayerState
).hidden.sum()
"
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0x57c9594]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0x57c702e]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0x57c977c]
/lib/x86_64-linux-gnu/libpthread.so.0(+0x128a0)[0x7f9099b458a0]
/lib/x86_64-linux-gnu/libc.so.6(gsignal+0xc7)[0x7f90983e7f47]
/lib/x86_64-linux-gnu/libc.so.6(abort+0x141)[0x7f90983e98b1]
/lib/x86_64-linux-gnu/libc.so.6(+0x3042a)[0x7f90983d942a]
/lib/x86_64-linux-gnu/libc.so.6(+0x304a2)[0x7f90983d94a2]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xfa910e]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xfa91cf]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xf4c484]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xf52431]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xf4ddd3]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xf4b206]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xaf005f]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xaf27d1]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xaecd78]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xaecd2b]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xb09c5a]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xaf7cb4]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xaecf45]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0xaf9c81]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0x697e79]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0x50d8ed]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0x50d22c]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0x500d1c]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0x49c103]
/lib/x86_64-linux-gnu/libc.so.6(__libc_start_main+0xe7)[0x7f90983cab97]
/swift-tensorflow-toolchain/usr/bin/swift-frontend[0x49bc1a]

To reproduce, I believe this should be reliably triggered by cloning tensorflow/swift-apis and building it using a current Swift.org toolchain via CMake:

cmake -B out -G Ninja -S swift-apis -D CMAKE_BUILD_TYPE=Release
cmake --build out