code: def __init__(self, config):
      self.config = config
      optimizer = config["optimizer"]
  
      def eval(state, ctx, tgt, ctx_length):
          def eval_loss(x, y, mask):
              transformer = CausalTransformerShard(config)
              return transformer.loss(x, y, mask=mask)
  
          eval_loss_fn = hk.without_apply_rng(
              hk.transform(eval_loss)
          ).apply
  
          mask = (jnp.arange(0, len(ctx)) > ctx_length) * -1e10
  
          return eval_loss_fn(to_bf16(state["params"]), ctx, tgt, mask)
  
      def train(state, ctx, tgt):
          def train_loss(x, y):
              transformer = CausalTransformerShard(config)
              out = transformer.loss(x, y, z_loss=True)
  
              return out["loss"], out["last_loss"]
  
          train_loss_fn = hk.without_apply_rng(
              hk.transform(train_loss)
          ).apply
  
          def microbatch(old_grad, batch):
              ctx, tgt = batch
  
              val_grad_fn = jax.value_and_grad(
                  train_loss_fn, has_aux=True
              )
              (loss, last_loss), grad = val_grad_fn(
                  to_bf16(state["params"]), ctx, tgt
              )
  
              new_grad = jax.tree_multimap(
                  lambda a, b: a + b, old_grad, grad
              )
              gnorm = global_norm(grad)
              return new_grad, (loss, last_loss, gnorm)
  
          if ctx.shape[0] == 1:
              val_grad_fn = jax.value_and_grad(
                  train_loss_fn, has_aux=True
              )
              (loss, last_loss), grad = val_grad_fn(
                  to_bf16(state["params"]), ctx[0], tgt[0]
              )
              gnorm = global_norm(grad)
          else:
              grad, (loss, last_loss, gnorm) = jax.lax.scan(
                  microbatch,
                  jax.tree_map(
                      lambda x: jnp.zeros_like(x).astype(jnp.bfloat16),
                      state["params"],
                  ),
                  (ctx, tgt),
              )
  
          grad_norm_micro = jax.lax.pmean(gnorm, "batch")
  
          grad = jax.lax.pmean(grad, "batch")
          grad_norm = global_norm(grad)
          updates, new_opt_state = optimizer.update(
              grad, state["opt_state"], state["params"]
          )
  
          return (
              to_f32(loss),
              to_f32(last_loss),
              to_f32(grad_norm),
              to_f32(grad_norm_micro),
              {
                  "params": optax.apply_updates(
                      state["params"], to_f32(updates)
                  ),
                  "step": state["step"] + 1,
                  "opt_state": new_opt_state,
              },
          )
  
      def init(key, x):
          def train_loss(x, y):
              transformer = CausalTransformerShard(config)
              return transformer.loss(x, y)
  
          param_init_fn = hk.transform(
              hk.experimental.optimize_rng_use(train_loss)
          ).init
  
          params = param_init_fn(key, x, x)
  
          return {
              "params": ("early_cast" in config and to_bf16 or to_f32)(
                  params
              ),
              "step": np.array(0),
              "opt_state": optimizer.init(params),
          }
  
      def generate(state, key, ctx, ctx_length, aux, sampler_options):
          sampler = config["sampler"]
          gen_length = self.gen_length
  
          def generate_sample(context, ctx_length, aux):
              transformer = CausalTransformerShard(config)
              _, initial_state = transformer.generate_initial(
                  context, ctx_length
              )
  
              def generate_scan_fn(carry, sampler_input):
                  next_token, decode_state, sample_key = carry
                  sample_key, new_key = jax.random.split(sample_key)
  
                  logits, new_state = transformer.generate_once(
                      next_token, decode_state
                  )
                  next_token, sample_info = sampler(
                      sample_key,
                      logits,
                      sampler_input,
                      **sampler_options,
                  )
  
                  if self.return_logits:
                      output = (next_token, sample_info, logits)
                  else:
                      output = (next_token, sample_info)
                  new_carry = (next_token, new_state, new_key)
                  return new_carry, output
  
              final_state, outputs = jax.lax.scan(
                  generate_scan_fn,
                  initial_state,
                  xs=aux,
                  length=gen_length,
              )
              return final_state, outputs
  
          generate_fn = hk.transform(generate_sample).apply
          return generate_fn(state["params"], key, ctx, ctx_length, aux)
  
      self.init_xmap = jax.experimental.maps.xmap(
          fun=init,
          in_axes=(["shard", ...], ["batch", ...]),
          out_axes=["shard", ...],
          axis_resources={"shard": "mp", "batch": "dp"},
      )
  
      self.eval_xmap = jax.experimental.maps.xmap(
          fun=eval,
          in_axes=(
              ["shard", ...],
              ["batch", ...],
              ["batch", ...],
              ["batch", ...],
          ),
          out_axes=["batch", ...],
          axis_resources={"shard": "mp", "batch": "dp"},
      )
  
      self.train_xmap = jax.experimental.maps.xmap(
          fun=train,
          in_axes=(["shard", ...], ["batch", ...], ["batch", ...]),
          out_axes=(
              ["batch", ...],
              ["batch", ...],
              ["batch", ...],
              ["batch", ...],
              ["shard", ...],
          ),
          donate_argnums=(0,),
          axis_resources={"shard": "mp", "batch": "dp"},
      )
  
      self.generate_xmap = jax.experimental.maps.xmap(
          fun=generate,
          in_axes=(
              ["shard", ...],
              ["batch", ...],
              ["batch", ...],
              ["batch", ...],
              ["batch", ...],
              ["batch", ...],
          ),
          out_axes=["batch", ...],
          axis_resources={"shard": "mp", "batch": "dp"},
      )
  
      self.move_xmap = jax.experimental.maps.xmap(
          fun=lambda x, _: to_bf16(x),
          in_axes=(["shard", ...], ["batch", ...]),
          out_axes=["shard", ...],
          axis_resources={"shard": "mp", "batch": "dp"},
      )
  
      key = hk.PRNGSequence(42)
  
      assert (
          thread_resources.env.shape["mp"]
          == config["cores_per_replica"]
      )
  
      dp = thread_resources.env.shape["dp"]
      mp = thread_resources.env.shape["mp"]
  
      mp_per_host = min(mp, 8)
  
      seq = config["seq"]
      vocab = config["n_vocab"]
  
      example_shape = (
          max(dp // jax.host_count(), 1),
          seq,
      )
      x = jax.random.uniform(
          next(key), example_shape, minval=0, maxval=vocab
      ).astype(
          jnp.uint32
      )  # batch, len
  
      head_print("key shape", jnp.array(key.take(mp_per_host)).shape)
      head_print("in shape", x.shape)
  
      head_print("dp", dp)
      head_print("mp", mp)
  
      self.gen_length = 1
      self.state = self.init_xmap(jnp.array(key.take(mp_per_host)), x)
  
      param_count = hk.data_structures.tree_size(self.state["params"])
      head_print(f"Total parameters: {param_count}")
  
imports: import gc
  import random
  import time
  from functools import partial
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  from jax.experimental.maps import thread_resources
  from jax.experimental.pjit import pjit
  from mesh_transformer.checkpoint import read_ckpt, write_ckpt, write_ckpt_v2, load_ckpt_v2
  from mesh_transformer.layers import EmbeddingShard, TransformerLayerShard, RelativePositionEmbs, ProjectionShard, TransformerLayerShardV2, Projection, EmbeddingShardV2
  from mesh_transformer.util import to_f32, to_bf16, maybe_shard, head_print, global_norm
  from jax.experimental import PartitionSpec
label_desc: Use image node_id 1 to create a new CausalTransformer object  with example: obj = CausalTransformer(config)