Sample one entry matching (block_pos, timestep_index) when 2D, or just timestep_index when 1D. Non-owned timestep buckets are transparently remapped to the nearest owned one. Returns None if the (remapped) bucket is empty.
(self, timestep_index, device, dtype, block_pos=None)
| 143 | |
| 144 | # ------------------------------------------------------------------ sample |
| 145 | def sample(self, timestep_index, device, dtype, block_pos=None): |
| 146 | """Sample one entry matching (block_pos, timestep_index) when 2D, or |
| 147 | just timestep_index when 1D. Non-owned timestep buckets are |
| 148 | transparently remapped to the nearest owned one. Returns None if the |
| 149 | (remapped) bucket is empty.""" |
| 150 | t = self._nearest_owned_t(self._t_bucket(timestep_index)) |
| 151 | key = self._make_key(t, block_pos) |
| 152 | buf = self.buckets[key] |
| 153 | if not buf: |
| 154 | return None |
| 155 | err = random.choice(buf) |
| 156 | return self._modulate(err).to(device=device, dtype=dtype) |
| 157 | |
| 158 | def sample_pos_any_t(self, block_pos, device, dtype): |
| 159 | """For 2D buffers: sample at the given position, with random timestep. |
nothing calls this directly
no test coverage detected