add helper for flax to be omnistaging-compatible (#4004)
add helper for flax to be omnistaging-compatible (#4004)


google/jax,google/jax,tensorflow/probability,google/jax,tensorflow/probability,google/jax