class_docstr: Per device, samples an item from iterator and places on device.
  
  if pmapped_user:
    Items from the resulting generator are intended to be used in a pmapped
    function. Every element is a ShardedDeviceArray or (nested) Python container
    thereof. A single next() call to this iterator results in len(devices)
    calls to the underlying iterator. The returned items are put one on each
    device.
  if not pmapped_user:
    Places a sample from the iterator on the given device.
  
  Yields:
    If no split_fn is specified:
      DeviceArray/ShardedDeviceArray or (nested) Python container thereof
      representing the elements of shards stacked together, with each shard
      backed by physical device memory specified by the corresponding entry in
      devices.
  
    If split_fn is specified:
      PrefetchingSplit where the .host element is a stacked numpy array or
      (nested) Python contained thereof. The .device element is a
      DeviceArray/ShardedDeviceArray or (nested) Python container thereof.
  
  Raises:
    StopIteration: if there are not enough items left in the iterator to place
      one sample on each device.
    Any error thrown by the iterable_function. Note this is not raised inside
      the producer, but after it finishes executing.
class_docstr_tok: ["Per","device",",","samples","an","item","from","iterator","and","places","on","device",".","if","pmapped_user",":","Items","from","the","resulting","generator","are","intended","to","be","used","in","a","pmapped","function",".","Every","element","is","a","ShardedDeviceArray","or","(","nested",")","Python","container","thereof",".","A","single","next","(",")","call","to","this","iterator","results","in","len","(","devices",")","calls","to","the","underlying","iterator",".","The","returned","items","are","put","one","on","each","device",".","if","not","pmapped_user",":","Places","a","sample","from","the","iterator","on","the","given","device",".","Yields",":","If","no","split_fn","is","specified",":","DeviceArray\/ShardedDeviceArray","or","(","nested",")","Python","container","thereof","representing","the","elements","of","shards","stacked","together",",","with","each","shard","backed","by","physical","device","memory","specified","by","the","corresponding","entry","in","devices",".","If","split_fn","is","specified",":","PrefetchingSplit","where","the",".host","element","is","a","stacked","numpy","array","or","(","nested",")","Python","contained","thereof",".","The",".device","element","is","a","DeviceArray\/ShardedDeviceArray","or","(","nested",")","Python","container","thereof",".","Raises",":","StopIteration",":","if","there","are","not","enough","items","left","in","the","iterator","to","place","one","sample","on","each","device",".","Any","error","thrown","by","the","iterable_function",".","Note","this","is","not","raised","inside","the","producer",",","but","after","it","finishes","executing","."]
docstr: Constructs PutToDevicesIterable.
  
  Args:
    iterable: A python iterable. This is used to build the python prefetcher.
      Note that each iterable should only be passed to this function once as
      iterables aren't thread safe.
    pmapped_user: whether the user of data from this iterator is implemented
      using pmapping.
    devices: Devices used for prefecthing.
    split_fn: Optional function applied to every element from the iterable to
      split the parts of it that will be kept in the host and the parts that
      will sent to the device.
  
  Raises:
    ValueError: If devices list is empty, or if pmapped_use=False and more
      than 1 device is provided.
code: def __init__(
      self,
      iterable: Iterable[types.NestedArray],
      pmapped_user: bool,
      devices: Sequence[jax.Device],
      split_fn: Optional[_SplitFunction] = None,
  ):
      """Constructs PutToDevicesIterable.
  
      Args:
        iterable: A python iterable. This is used to build the python prefetcher.
          Note that each iterable should only be passed to this function once as
          iterables aren't thread safe.
        pmapped_user: whether the user of data from this iterator is implemented
          using pmapping.
        devices: Devices used for prefecthing.
        split_fn: Optional function applied to every element from the iterable to
          split the parts of it that will be kept in the host and the parts that
          will sent to the device.
  
      Raises:
        ValueError: If devices list is empty, or if pmapped_use=False and more
          than 1 device is provided.
      """
      self.num_devices = len(devices)
      if self.num_devices == 0:
          raise ValueError("At least one device must be specified.")
      if (not pmapped_user) and (self.num_devices != 1):
          raise ValueError(
              "User is not implemented with pmapping but len(devices) "
              f"= {len(devices)} is not equal to 1! Devices given are:"
              f"\n{devices}"
          )
  
      self.iterable = iterable
      self.pmapped_user = pmapped_user
      self.split_fn = split_fn
      self.devices = devices
      self.iterator = iter(self.iterable)
  
imports: import functools
  import itertools
  import queue
  import threading
  from typing import Callable, Iterable, Iterator, NamedTuple, Optional, Sequence, Tuple, TypeVar
  from absl import logging
  from acme import core
  from acme import types
  from acme.jax import types
  import haiku
  import jax
  import jax.numpy
  import numpy
  import reverb
  import tree
label_desc: Use image node_id 1 to create a new PutToDevicesIterable object from inherited base classes: Iterable with example: obj = PutToDevicesIterable(iterable, pmapped_user, devices, split_fn)