def _from_bits(self, bits):
    """Converts a list of bit-cast weights back to weights."""
    # This is the reverse of _to_bits, see above for explanation.
    if not fastmath.is_backend(fastmath.Backend.JAX):
      return bits
    weights = []
    for b in bits:
      if b.dtype == np.uint16:  # currently all uint16 are bfloat16s
        w = jax.lax.bitcast_convert_type(b, jnp.bfloat16)
        weights.append(np.asarray(w))
      else:
        weights.append(b)
    return weights
Converts a list of bit-cast weights back to weights.
trax/supervised/training.py