(self)
| 135 | self.assertEqual(first_token.log_prob.shape, (1, 1)) |
| 136 | |
| 137 | def test_basic_decode(self): |
| 138 | devices_array = maxtext_utils.create_device_mesh(self.cfg) |
| 139 | mesh = Mesh(devices_array, self.cfg.mesh_axes) |
| 140 | quant = quantizations.configure_quantization(self.cfg) |
| 141 | model = models.transformer_as_linen(config=self.cfg, mesh=mesh, quant=quant, model_mode=MODEL_MODE_PREFILL) |
| 142 | ids, decoder_segment_ids, decoder_positions = self.get_data() |
| 143 | |
| 144 | transformer_vars = model.init( |
| 145 | {"params": self.rng, "aqt": self.rng, "dropout": self.rng}, |
| 146 | ids, |
| 147 | decoder_positions, |
| 148 | decoder_segment_ids, |
| 149 | enable_dropout=False, |
| 150 | ) |
| 151 | input_tokens = jnp.array([1, 306, 5360, 304]) |
| 152 | engine = MaxEngine(self.cfg, jax.devices()) |
| 153 | params = engine.load_params(params=transformer_vars) |
| 154 | decode_state = engine.init_decode_state() |
| 155 | prefill_result, _ = engine.prefill(params=params, padded_tokens=input_tokens, true_length=4) |
| 156 | decode_state = engine.insert(prefill_result, decode_state, slot=0) |
| 157 | decode_state, result_token = engine.generate(params=params, decode_state=decode_state) |
| 158 | |
| 159 | self.assertEqual(result_token.log_prob.ndim, 2) |
| 160 | self.assertEqual(result_token.log_prob.shape[1], 1) |
| 161 | self.assertEqual(result_token.data.ndim, 2) |
| 162 | self.assertEqual(result_token.data.shape[1], 3) |
| 163 | |
| 164 | @pytest.mark.skip(reason="Can only pass on CPU.") |
| 165 | def test_chunked_prefill(self): |
nothing calls this directly
no test coverage detected