keras/src/ops/core.py/Switch/compute_output_spec
class Switch:    def compute_output_spec(self, index, branches, *operands):
        # We use first branch for output_spec
        spec = backend.compute_output_spec(branches[0], *operands)
        return spec