rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator Class Reference#
rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator Class Reference
Inheritance diagram for rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator:
Public Member Functions | |
| def | __init__ (self, pipelines, sharding=None) |
| def | set_element_spec (self, outputs) |
| def | assert_output_shape_and_type (self, outputs) |
| def | next_with_peek_impl (self) |
| def | __next__ (self) |
| def | __iter__ (self) |
| def | peek (self) |
| def | peek_async (self) |
| def | element_spec (self) |
 Public Member Functions inherited from rocAL_pybind.amd.rocal.plugin.jax.ROCALJaxIterator | |
| def | next (self) |
| def | reset (self) |
| def | place_output_with_device_put (self, individual_outputs) |
| def | place_output_with_sharding (self, individual_outputs) |
| def | __len__ (self) |
| def | __del__ (self) |
Data Fields | |
| mutex | |
| pool | |
| peek | |
| element_spec | |
| Data Fields inherited from rocAL_pybind.amd.rocal.plugin.jax.ROCALJaxIterator | |
| pipelines | |
| num_devices | |
| batch_size | |
| iterator_length | |
| last_batch_policy | |
| sharding | |
Detailed Description
ROCALJaxIterator extended with peek functionality. Compatible with Google CLU PeekableIterator. Reference: https://github.com/google/CommonLoopUtils/blob/main/clu/data/dataset_iterator.py
Member Function Documentation
◆ __iter__()
| def rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator.__iter__ | ( | self | ) |
Returns the iterator object. Resets if the iterator has been used.
Reimplemented from rocAL_pybind.amd.rocal.plugin.jax.ROCALJaxIterator.
◆ __next__()
| def rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator.__next__ | ( | self | ) |
Returns the next element from the iterator and advances it.
Reimplemented from rocAL_pybind.amd.rocal.plugin.jax.ROCALJaxIterator.
◆ assert_output_shape_and_type()
| def rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator.assert_output_shape_and_type | ( | self, | |
| outputs | |||
| ) |
Asserts that the shape and type of the outputs are consistent.
Args:
outputs: The output from the iterator.
Raises:
ValueError: If the shape or type of the output changes between iterations.
Returns:
The outputs if they are consistent.
◆ element_spec()
| def rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator.element_spec | ( | self | ) |
Returns the element spec for the elements returned by the iterator.
◆ next_with_peek_impl()
| def rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator.next_with_peek_impl | ( | self | ) |
Returns the next element from the iterator and advances the iterator.
◆ peek()
| def rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator.peek | ( | self | ) |
Returns the next element from the iterator without advancing the iterator.
◆ peek_async()
| def rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator.peek_async | ( | self | ) |
Returns future that will return the next element from the iterator without advancing the iterator.
◆ set_element_spec()
| def rocAL_pybind.amd.rocal.plugin.jax.ROCALPeekableIterator.set_element_spec | ( | self, | |
| outputs | |||
| ) |
Sets the element spec for the iterator.
Args:
outputs: The output from the iterator.
The documentation for this class was generated from the following file:
- /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-rocal/checkouts/develop/rocAL_pybind/amd/rocal/plugin/jax.py