[flax] Add messages to some assertions in pool to help with debugging.#1321

copybara-service[bot] merged 1 commit into

mastergoogle/flax:masterfrom

test_374084282google/flax:test_374084282Copy head branch name to clipboard