code: def _fold_in(rng, d):
      """Equivalent of jax.random.fold_in."""
      # TODO(lukaszkaiser): verify that this function has good randomness
      # properties or switch to an implementation equivalent to JAX.
      _, rng = tf_np_extensions.split(
          rng + tf_np.sum(d).astype(tf_np.int64), 2
      )
      return rng
  
file_path: trax/trax/fastmath/tf.py
imports: import numpy
  import tensorflow.compat.v2
  from trax.shapes import ShapeDtype
  from trax.tf_numpy import extensions
  from trax.tf_numpy import numpy
label_desc: Use image node_id 8 for calling a global function with example usage: _fold_in(rng, d) and returns: rng