class_docstr: Accelerates a layer, running in data-parallel way on multiple devices.
  
  By default it uses all available accelerators, splits the input on the
  first (batch) axis, and runs each part on the corresponding accelerator.
  If only one accelerator is available, this layer JIT-compiles the underlying
  layer and in this way makes it run faster.
  
  The output is guaranteed to be the same as the output of the original layer
  if the batch dimension is divisible by the number of devices. If it is not,
  then 0-padding is added to make it divisible and the output may be affected
  if it relies on layers like batch normalization.
  
  This layer does not require calling ``init`` if the underlying layer has
  already been initialized, so it can be used as follows::
  
      layer = tl.Serial(...)
      layer.init(...)
      fast_layer = tl.Accelerate(layer)
      y = fast_layer(x)  # Split x on batch and run data-parallel
  
  In case the weights of this layer need to be set using the weights of
  the sublayer, use the ``replicate_weights`` function::
  
      # Instead of layer.weights = new_weights:
      fast_layer.replicate_weights(new_weights)
class_docstr_tok: ["Accelerates","a","layer",",","running","in","data-parallel","way","on","multiple","devices",".","By","default","it","uses","all","available","accelerators",",","splits","the","input","on","the","first","(","batch",")","axis",",","and","runs","each","part","on","the","corresponding","accelerator",".","If","only","one","accelerator","is","available",",","this","layer","JIT-compiles","the","underlying","layer","and","in","this","way","makes","it","run","faster",".","The","output","is","guaranteed","to","be","the","same","as","the","output","of","the","original","layer","if","the","batch","dimension","is","divisible","by","the","number","of","devices",".","If","it","is","not",",","then","0-padding","is","added","to","make","it","divisible","and","the","output","may","be","affected","if","it","relies","on","layers","like","batch","normalization",".","This","layer","does","not","require","calling","``","init","``","if","the","underlying","layer","has","already","been","initialized",",","so","it","can","be","used","as","follows",":",":","layer","=","tl.Serial","(","...",")","layer.init","(","...",")","fast_layer","=","tl.Accelerate","(","layer",")","y","=","fast_layer","(","x",")","#","Split","x","on","batch","and","run","data-parallel","In","case","the","weights","of","this","layer","need","to","be","set","using","the","weights","of","the","sublayer",",","use","the","``","replicate_weights","``","function",":",":","#","Instead","of","layer.weights","=","new_weights",":","fast_layer.replicate_weights","(","new_weights",")"]
docstr: Sets the weights of the sublayer and replicates them for this layer.
docstr_tok: ["Sets","the","weights","of","the","sublayer","and","replicates","them","for","this","layer","."]
code: def replicate_weights(self, weights):
      """Sets the weights of the sublayer and replicates them for this layer."""
      self.sublayer.weights = weights
      self._weights = self._prepare_weights(weights)
  
imports: import jax
  import numpy
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
label_desc: Use image node_id 6 for calling the Accelerate obj's underlying member method code with example usage: obj.replicate_weights(weights) without return types