(self)
| 67 | return config |
| 68 | |
| 69 | def get_data(self): |
| 70 | s = (self.cfg.global_batch_size_to_train_on, self.cfg.max_target_length) |
| 71 | ids = jax.random.randint(self.rng, s, 0, self.cfg.vocab_size) |
| 72 | |
| 73 | decoder_segment_ids = jax.numpy.zeros(s) + DECODING_ACTIVE_SEQUENCE_INDICATOR |
| 74 | decoder_positions = jnp.stack( |
| 75 | [jnp.arange(self.cfg.max_target_length, dtype=jnp.int32) for _ in range(self.cfg.global_batch_size_to_train_on)] |
| 76 | ) |
| 77 | |
| 78 | return ids, decoder_segment_ids, decoder_positions |
| 79 | |
| 80 | def test_stack_and_unstack_prefill_cache(self): |
| 81 | config = pyconfig.initialize( |
no outgoing calls
no test coverage detected