[JAX] Support Python 3 in JAX.

Notable changes:
* __builtins__ -> six.moves.builtins
* xrange -> six.moves.xrange
* reduce -> six.moves.reduce
* map/filter return generators, so add explicit list() calls around many to avoid a semantic change. Probably too conservative, but it's easier to add them everywhere than work out which are strictly necessary.
* no long in Python 3
* raise TypeError, x -> raise TypeError(x)
* gen.next() -> next(gen)
* string / bytes
* no tuple destructuring in lambda arguments.
* / and np.divide changed semantics, use // and np.floor_divide.

PiperOrigin-RevId: 222451919