Merge pull request jax-ml#25179 from rajasekharporeddy:lax_switch

PiperOrigin-RevId: 702018052

Author: Google-ML-Automation

jax/_src/lax/control_flow/conditionals.py

@@ -87,6 +87,7 @@ def switch(index, branches, *operands):
   Args:
     index: Integer scalar type, indicating which branch function to apply.
     branches: Sequence of functions (A -> B) to be applied based on ``index``.
+      All branches must return the same output structure.
     operands: Operands (A) input to whichever branch is applied.
 
   Returns: