Repository: google/flax

Issue:
Clarification regarding LM1B input pipeline
Hi, I am looking for two clarifications regarding the [input_pipeline](https://github.com/google/flax/blob/master/examples/lm1b/input_pipeline.py) in Flax LM1B example. 
1. I think there might be a bug at 
   https://github.com/google/flax/blob/master/examples/lm1b/input_pipeline.py#L132. 
   ```
   max_eval_length = max_eval_length or target_bucket_length * 32
   bucket_boundaries[-1] = max_eval_length
   bucket_batch_sizes[-1] = target_batch_size // max_eval_length
   ```
   The last statement might result in 0 batch size for the last bucket. 

   If  `max_eval_length == target_bucket_length * 32`, the `bucket_batch_size[-1]` should be `target_batch_size // 32` instead of `target_batch_size // (32 * target_bucket_length)` (which is what current implementation does). In general, `max_eval_length >> target_batch_size`, hence, this might result in 0 batch size for the last bucket. 

2. The documentation [here](https://github.com/google/flax/blob/master/examples/lm1b/input_pipeline.py#L241) mentions that 
dynamic batching is currently not compatible with multiple hosts, although the bucketing function handles the case when `n_devices > 1`. Currently, if I understand the control flow correctly, the binning (and batching) happens first, then followed by distribution across the hosts through pmap. If this is the case and the effective batch size is ensured to be a multiple of `n_devices`, why should dynamic batching be any different from the static batching control flow? 



Discussion:
1. This particular function is more than a bit convoluted!  My apologies for the unclear logic, we'll probably just remove this special case altogether - we're just trying to deal w. the more general case of evaluating on eval-set examples much longer than those in the training set, which I've needed to do occasionally but haven't used in ages!  
This does look incorrect, I'll make a fix, I believe the correct expression to attempt to maintain the same total per-batch token count is:  
`bucket_batch_sizes[-1] = target_batch_size // (max_eval_length // target_bucket_length)`
However, this doesn't cause much trouble, this never divides by zero, since that's forced to be 1 or more at: https://github.com/google/flax/blob/master/examples/lm1b/input_pipeline.py#L137  for long eval examples this usually ends up having a batch size of 1 anyway, which is why we didn't notice this earlier.  Thanks for pointing it out!

2.  multiple __hosts__ are not the same thing as multiple __devices__ : this function works fine for multiple devices, but in the case of multiple hosts each with their own set of devices, we would need to synchronize the "bucket" being sharded and fed to the devices on each host - in JAX the multihost programming model requires each host to feed its own devices with exactly the same input shape at each synchronized pmap step.

Resolving patch:
diff --git a/examples/lm1b/input_pipeline.py b/examples/lm1b/input_pipeline.py
--- a/examples/lm1b/input_pipeline.py
+++ b/examples/lm1b/input_pipeline.py
@@ -129,7 +129,8 @@ def bin_and_batch(dataset,
     if not training:
       max_eval_length = max_eval_length or target_bucket_length * 32
       bucket_boundaries[-1] = max_eval_length
-      bucket_batch_sizes[-1] = target_batch_size // max_eval_length
+      bucket_batch_sizes[-1] = (target_batch_size // 
+                                (max_eval_length // target_bucket_length))
     # We will pad to boundaries which pads to bucket_boundary-1: add 1 here.
     bucket_boundaries = [b + 1 for b in bucket_boundaries]
     # Make batch sizes divisible by n_devices.


Test patch: