Skip to content

Improve RTD contents and structure. #1198

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Jan 16, 2026
Merged

Improve RTD contents and structure. #1198

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/data_sources/protocol.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,5 +113,5 @@ File systems
Grain supports the formats mentioned above in combination with the following
file systems (in addition to the local file system):

* :doc:`S3 <../tutorials/data_sources/load_from_s3_tutorial>`
* :doc:`GCS <../tutorials/data_sources/load_from_gcs_tutorial>`
* :doc:`Amazon S3 <../tutorials/data_sources/load_from_s3_tutorial>`
13 changes: 8 additions & 5 deletions docs/grain.checkpoint.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,16 @@
===========================

.. automodule:: grain.checkpoint
.. currentmodule:: grain.checkpoint

List of Members
---------------

.. autosummary::
:toctree: _autosummary
.. autoclass:: CheckpointHandler
:members:

CheckpointHandler
CheckpointSave
CheckpointRestore
.. autoclass:: CheckpointSave
:members:

.. autoclass:: CheckpointRestore
:members:
22 changes: 13 additions & 9 deletions docs/grain.constants.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,21 @@
==========================

.. automodule:: grain.constants
:noindex:

List of Constants
-----------------

.. autosummary::
:toctree: _autosummary
.. autodata:: grain.constants.DATASET_INDEX

DATASET_INDEX
EPOCH
INDEX
META_FEATURES
RECORD
RECORD_KEY
SEED
.. autodata:: grain.constants.EPOCH

.. autodata:: grain.constants.INDEX

.. autodata:: grain.constants.META_FEATURES

.. autodata:: grain.constants.RECORD

.. autodata:: grain.constants.RECORD_KEY

.. autodata:: grain.constants.SEED
28 changes: 28 additions & 0 deletions docs/grain.data_loader.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
``grain`` DataLoader
=================

.. automodule:: grain._src.python.data_loader
.. currentmodule:: grain

List of Members
---------------

.. autofunction:: load

.. autoclass:: DataLoader
:special-members: __init__, __iter__
:members:
:undoc-members:

.. autoclass:: DataLoaderIterator
:special-members: __init__, __iter__, __next__
:members:
:undoc-members:

.. autoclass:: Record
:members:
:undoc-members:

.. autoclass:: RecordMetadata
:members:
:undoc-members:
2 changes: 2 additions & 0 deletions docs/grain.dataset.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,3 +33,5 @@ List of Members
:show-inheritance:
:inherited-members:
:undoc-members:

.. autoclass:: ReadOptions
140 changes: 100 additions & 40 deletions docs/grain.experimental.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,43 +6,103 @@
List of Members
---------------

.. autosummary::
:toctree: _autosummary

FlatMapTransform
DatasetOptions
ExecutionTrackingMode
apply_transformations
ElasticIterator
WithOptionsIterDataset
ParquetIterDataset
TFRecordIterDataset
batch_and_pad
CacheIterDataset
FlatMapMapDataset
FlatMapIterDataset
InterleaveIterDataset
LimitIterDataset
RngPool
FirstFitPackIterDataset
BestFitPackIterDataset
BOSHandling
ConcatThenSplitIterDataset
multithread_prefetch
ThreadPrefetchIterDataset
ThreadPrefetchDatasetIterator
RebatchIterDataset
RepeatIterDataset
WindowShuffleMapDataset
WindowShuffleIterDataset
ZipMapDataset
ZipIterDataset
PackAndBatchOperation
index_shuffle
assert_equal_output_after_checkpoint
device_put
PerformanceConfig
pick_performance_config
get_element_spec
set_next_index
get_next_index
.. autoclass:: FlatMapTransform
:members:

.. autoclass:: DatasetOptions
:members:

.. autoclass:: ExecutionTrackingMode
:members:

.. autofunction:: apply_transformations

.. autoclass:: ElasticIterator
:special-members: __init__, __iter__, __next__

.. autoclass:: WithOptionsIterDataset
:special-members: __init__, __iter__

.. autoclass:: ParquetIterDataset
:special-members: __init__, __iter__

.. autoclass:: TFRecordIterDataset
:special-members: __init__, __iter__

.. autofunction:: batch_and_pad

.. autoclass:: CacheIterDataset
:special-members: __init__, __iter__

.. autoclass:: FlatMapMapDataset
:special-members: __init__, __getitem__

.. autoclass:: FlatMapIterDataset
:special-members: __init__, __iter__

.. autoclass:: InterleaveIterDataset
:special-members: __init__, __iter__

.. autoclass:: LimitIterDataset
:special-members: __init__, __iter__

.. autoclass:: RngPool
:members:

.. autoclass:: FirstFitPackIterDataset
:special-members: __init__, __iter__

.. autoclass:: BestFitPackIterDataset
:special-members: __init__, __iter__

.. autoclass:: BOSHandling
:members:

.. autoclass:: ConcatThenSplitIterDataset
:special-members: __init__, __iter__

.. autofunction:: multithread_prefetch

.. autoclass:: ThreadPrefetchIterDataset
:special-members: __init__, __iter__

.. autoclass:: ThreadPrefetchDatasetIterator
:special-members: __init__, __iter__, __next__

.. autoclass:: RebatchIterDataset
:special-members: __init__, __iter__

.. autoclass:: RepeatIterDataset
:special-members: __init__, __iter__


.. autoclass:: WindowShuffleMapDataset
:special-members: __init__, __getitem__

.. autoclass:: WindowShuffleIterDataset
:special-members: __init__, __iter__


.. autoclass:: ZipMapDataset
:special-members: __init__, __getitem__


.. autoclass:: ZipIterDataset
:special-members: __init__, __iter__

.. autofunction:: index_shuffle

.. autofunction:: assert_equal_output_after_checkpoint

.. autofunction:: device_put

.. autoclass:: PerformanceConfig
:members:

.. autofunction:: pick_performance_config

.. autofunction:: get_element_spec

.. autofunction:: set_next_index

.. autofunction:: get_next_index
8 changes: 4 additions & 4 deletions docs/grain.multiprocessing.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
List of Members
---------------

.. autosummary::
:toctree: _autosummary
.. autoclass:: MultiprocessingOptions

.. autoclass:: SharedMemoryArray
:members:

MultiprocessingOptions
SharedMemoryArray
10 changes: 8 additions & 2 deletions docs/grain.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,15 @@
.. currentmodule:: grain

Public API: ``grain`` package
``grain`` package
=============================

.. toctree::
:hidden:
:caption: Core APIs

Dataset APIs <grain.dataset>
DataLoader APIs <grain.data_loader>

Subpackages
-----------

Expand Down Expand Up @@ -35,7 +42,6 @@ Simple high-level pipelines


.. autosummary::
:toctree: _autosummary

load
DataLoader
Expand Down
12 changes: 7 additions & 5 deletions docs/grain.samplers.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,11 @@
List of Members
---------------

.. autosummary::
:toctree: _autosummary
.. autoclass:: IndexSampler
:special-members: __init__, __len__, __getitem__

IndexSampler
Sampler
SequentialSampler
.. autoclass:: Sampler
:special-members: __init__, __len__, __getitem__

.. autoclass:: SequentialSampler
:special-members: __init__, __len__, __getitem__
10 changes: 5 additions & 5 deletions docs/grain.sharding.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@
List of Members
---------------

.. autosummary::
:toctree: _autosummary
.. autoclass:: ShardOptions

.. autoclass:: NoSharding

.. autoclass:: ShardByJaxProcess

NoSharding
ShardByJaxProcess
ShardOptions
26 changes: 15 additions & 11 deletions docs/grain.transforms.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,18 @@
List of Members
---------------

.. autosummary::
:toctree: _autosummary

Batch
Filter
Map
MapWithIndex
RandomMap
Transformation
Transformations
DatasetSelectionMap
.. autoclass:: Batch

.. autoclass:: Filter

.. autoclass:: Map

.. autoclass:: MapWithIndex

.. autoclass:: RandomMap

.. autoclass:: Transformation

.. autoclass:: Transformations

.. autoclass:: DatasetSelectionMap
Loading
Loading