jaxpi.samplers
Infinite collocation-point samplers.
BaseSampler [source]
Infinite iterator over randomly generated batches.
The global batch is sharded across devices by the training step, so the batch size must be divisible by the number of devices.
BaseSampler.data_generation() [source]
BaseSampler.data_generation(self, key)UniformSampler [source]
Bases: BaseSampler
Infinite iterator over randomly generated batches.
The global batch is sharded across devices by the training step, so the batch size must be divisible by the number of devices.
UniformSampler.data_generation()
UniformSampler.data_generation(self, key)Generates data containing batch_size samples
MeshSampler [source]
Bases: BaseSampler
Infinite iterator over randomly generated batches.
The global batch is sharded across devices by the training step, so the batch size must be divisible by the number of devices.
MeshSampler.data_generation()
MeshSampler.data_generation(self, key)Generates data containing batch_size samples.
TemporalMeshSampler [source]
Bases: BaseSampler
Infinite iterator over randomly generated batches.
The global batch is sharded across devices by the training step, so the batch size must be divisible by the number of devices.
TemporalMeshSampler.data_generation()
TemporalMeshSampler.data_generation(self, key)Generates data containing batch_size samples