code: def _AsTokenIDs():
      """Returns a layer that makes mask values look like token ID ints."""
      return tl.Fn("AsTokenIDs", lambda x: x.astype(jnp.int32))
  
file_path: trax/trax/models/reformer/reformer.py
imports: from trax import layers
  from trax.fastmath import numpy
  from trax.models.research import configurable_transformer
label_desc: Use image node_id 9 for calling a global function with example usage: _AsTokenIDs() and returns: tl