My pyproject.toml
[project]
name = "test-grain"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
"grain>=0.2.18",
"jax>=0.6.0",
]
my script
"""Bare-bones repro of the synthpix grain pipeline: source -> to_iter_dataset
-> mp_prefetch -> map. Just random data, no synthpix/jax. Run with the
feature-new-rendering venv (grain 0.2.18)."""
import numpy as np
from absl import flags
from grain import MapDataset, MultiprocessingOptions, ReadOptions
def main():
flags.FLAGS.mark_as_parsed()
# random "flow fields" as the data source (a plain sequence of ndarrays)
data = [np.random.rand(2, 256, 256, 2).astype(np.float32) for _ in range(50)]
ds = MapDataset.source(data)
ds = ds.to_iter_dataset(ReadOptions(num_threads=1, prefetch_buffer_size=1))
ds = ds.mp_prefetch(MultiprocessingOptions(num_workers=2, per_worker_buffer_size=3))
ds = ds.map(lambda x: x)
print("building iterator...")
it = iter(ds)
try:
print("pulling one batch...")
batch = next(it)
# copy out of shm now, so its async unlink isn't racing interpreter exit
batch = np.array(batch, copy=True)
print("OK, got batch shape:", batch.shape)
finally:
it.close()
if __name__ == "__main__":
main()ISSUE: unless you use flags.FLAGS.mark_as_parsed(), it fails. this happens only if jax is also imported.
My pyproject.toml
[project]
name = "test-grain"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
"grain>=0.2.18",
"jax>=0.6.0",
]
my script
ISSUE: unless you use flags.FLAGS.mark_as_parsed(), it fails. this happens only if jax is also imported.