def _to_bits(self, weights):
    """Converts a list of weights to bit-cast weights and their types."""
    # This is currently needed to pickle bfloat16 arrays from JAX.
    # TODO(lukaszkaiser): remove once it is not needed (the following unit test
    #   checks it: training_test/test_restores_step_bfloat16).
    if not fastmath.is_backend(fastmath.Backend.JAX):
      return weights
    bits = []
    for w in weights:
      if w.dtype == jnp.bfloat16:
        converted = jax.lax.bitcast_convert_type(w, np.uint16)
        bits.append(np.asarray(converted.astype(np.uint16)))
      else:  # for non-bfloat16 weights, be compatible with earlier checkpoints
        bits.append(np.asarray(w))
    return bits
Converts a list of weights to bit-cast weights and their types.
trax/supervised/training.py