class_docstr: Multi-device accelerated trainer.
  
  Given an optimizer and a composite layer containing model+loss, this class
  creates a multi-device accelerated function with which it can compute one step
  of updates to the model's weights/state and the optimizer slots. By default
  it uses all available accelerators, via JIT compilation and parallel mapping.
  
  The optimizer and model must be initialized prior to use by this class.
  
  The key `one_step` function runs one forward-backward pass through the model,
  and returns the resulting loss value and updated optimizer statistics. As a
  side effect, the function also modifies the model weights and optimizer slots.
class_docstr_tok: ["Multi-device","accelerated","trainer",".","Given","an","optimizer","and","a","composite","layer","containing","model+loss",",","this","class","creates","a","multi-device","accelerated","function","with","which","it","can","compute","one","step","of","updates","to","the","model","'s","weights\/state","and","the","optimizer","slots",".","By","default","it","uses","all","available","accelerators",",","via","JIT","compilation","and","parallel","mapping",".","The","optimizer","and","model","must","be","initialized","prior","to","use","by","this","class",".","The","key","`","one_step","`","function","runs","one","forward-backward","pass","through","the","model",",","and","returns","the","resulting","loss","value","and","updated","optimizer","statistics",".","As","a","side","effect",",","the","function","also","modifies","the","model","weights","and","optimizer","slots","."]
code: def __init__(
      self, model_with_loss, optimizer, n_devices=None, adasum=False
  ):
      self._model_with_loss = model_with_loss
      self._optimizer = optimizer
      self._n_devices = n_devices or fastmath.local_device_count()
      self._adasum = adasum
  
      # optimizer slots and opt_params may need to be replicated
      self._slots, self._opt_params = tl.on_cpu(
          tl.for_n_devices(
              (self._optimizer.slots, self._optimizer.opt_params),
              self._n_devices,
          )
      )
  
      # accelerated version of model+loss to replicate weights and state
      self._accelerated_model_with_loss = tl.Accelerate(
          model_with_loss, n_devices=n_devices
      )
  
      # Signature:
      # (batch, weights, state, rng) -> ((loss, state), gradients)
      self._forward_and_backward_fn = fastmath.value_and_grad(
          model_with_loss.pure_fn,
          argnums=1,  # arg1 of pure_fn: weights
          has_aux=True,
      )  # return (loss, state), gradients
  
      # Signature:
      # (weights, slots), step, opt_params, batch, state, rng ->
      # (weights, slots), state, stats
      self._accelerated_update_fn = _accelerate_update_fn(
          self._forward_and_backward_fn,
          self._optimizer,
          n_devices=self._n_devices,
          accelerate=True,
          adasum=self._adasum,
      )
  
imports: from concurrent import futures
  import functools
  import os
  import time
  from absl import logging
  import jax
  import numpy
  import psutil
  from trax import fastmath
  from trax import layers
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
label_desc: Use image node_id 1 to create a new Trainer object  with example: obj = Trainer(model_with_loss, optimizer, n_devices, adasum)