flax

429033b3 - The CL introduces a new SowConfig which is used to configure which intermediate values are to be sowed.

The CL introduces a new SowConfig which is used to configure which intermediate values are to be sowed. TransformerIntermediates and BlockIntermediates are used to store the intermediate values of multiple decoding steps. PiperOrigin-RevId: 723053959