Two extensions to circular pipeline
Two extensions to circular pipeline

1. Allow num_microbatches == num_stages, but async circular transfer will be
disabled in this case.

2. Add an option where the repeated layers share the same weight. This can be
useful for token-level circular auto-regressive decoding.

PiperOrigin-RevId: 459863053

google/praxis,google/praxis