[JAX SC] Continuation of cl/836767096. Templatize `kCreateBuckets` for compile-time optimization.#541

copybara-service[bot] merged 1 commit intomainjax-ml/jax-tpu-embedding:mainfrom

test_826579723jax-ml/jax-tpu-embedding:test_826579723Copy head branch name to clipboard