jax

965ba054 - Rewrite `tf.aliasing_output` when wrapping the main StableHLO function

Rewrite `tf.aliasing_output` when wrapping the main StableHLO function When the original main function has tokens and the wrapped main function does not, there are fewer outputs in the wrapped main than the original main. This is problematic for `tf.aliasing_output`, which is an argument attribute that stores result indexes to which the argument can alias. This CL makes the wrapper main creation rewrite `tf.aliasing_output` according to the new result indexes. The newly added test verifies that the aliasing indexes are correct across all supported serialization versions. Confirmed that the test fails without changes in export.py (versions 6, 7, and 8 fail and 9 passes). PiperOrigin-RevId: 587237842