TransformerLM lm1b model accept float and out-of vocabulary token as input

[google]

TransformerLM lm1b model accept float and out-of vocabulary token as input

[levskaya]

Yeah, currently the model forcibly casts vocab data as int32 type. We could be more strict and issue a warning instead (there is a type check in the embedding layer).

Out of vocab is trickier - we can't raise exceptions easily (or at all) on accelerators, so out-of-bounds access silently clamps. There is discussion in JAX to instead emit NaNs for out-of-bounds access which would error-out harder.
eg.

jnp.take(jnp.array([[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0]]), 
         jnp.array([0,1,2,3,4]), 
         axis=0)
# DeviceArray([[1., 0., 0.],
#              [0., 1., 0.],
#              [0., 0., 1.],
#              [0., 0., 1.],
#              [0., 0., 1.]], dtype=float32)

Thanks for raising the issue vis-a-vis the updated TFDS API - this loader dates back to before the new api existed, I've filed a separate issue #138 to review and update this in all of our examples.

I'll investigate the final issue, but I somewhat doubt it makes much difference given that the tfds datasets are set up as a lazy streams, rather than actually loading much on construction.

[marcvanzee]

Closing due to inactivity. If anyone would like to work on this, feel free to re-open!