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","."]
docstr: Returns the accelerated composite model+loss for this instance.
docstr_tok: ["Returns","the","accelerated","composite","model+loss","for","this","instance","."]
code: def accelerated_model_with_loss(self):
      """Returns the accelerated composite model+loss for this instance."""
      return self._accelerated_model_with_loss
  
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 3 for calling the Trainer obj's underlying member method code with example usage: obj.accelerated_model_with_loss() and returns: self