Can't run `bert_vocab_from_dataset` without `TypeError: Tensor is unhashable` when import `trax` with `tensorflow`

[mihalt]

Just want to let you know about this issue in tensorflow