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","."]
code: def __next__(self) -> types.NestedArray:
      try:
          if not self.pmapped_user:
              item = next(self.iterator)
              if self.split_fn is None:
                  return jax.device_put(item, self.devices[0])
              item_split = self.split_fn(item)
              return PrefetchingSplit(
                  host=item_split.host,
                  device=jax.device_put(
                      item_split.device, self.devices[0]
                  ),
              )
  
          items = itertools.islice(self.iterator, self.num_devices)
          items = tuple(items)
          if len(items) < self.num_devices:
              raise StopIteration
          if self.split_fn is None:
              return jax.device_put_sharded(tuple(items), self.devices)
          else:
              # ((host: x1, device: y1), ..., (host: xN, device: yN)).
              items_split = (self.split_fn(item) for item in items)
              # (host: (x1, ..., xN), device: (y1, ..., yN)).
              split = tree.map_structure_up_to(
                  PrefetchingSplit(None, None),
                  lambda *x: x,
                  *items_split,
              )
  
              return PrefetchingSplit(
                  host=np.stack(split.host),
                  device=jax.device_put_sharded(
                      split.device, self.devices
                  ),
              )
  
      except StopIteration:
          raise
  
      except Exception:  # pylint: disable=broad-except
          logging.exception("Error for %s", self.iterable)
          raise
  
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 3 for calling the PutToDevicesIterable obj's underlying member method code with example usage: obj.__next__() and returns: PrefetchingSplit, jax, PrefetchingSplit, jax