(self)
| 350 | precision: Optional[Union[jax.lax.Precision, str]]=None |
| 351 | |
| 352 | def setup(self): |
| 353 | self.transformer = FlaxVideoLLaMAModule(self.config, dtype=self.dtype) |
| 354 | self.vision_head = nn.Dense( |
| 355 | self.config.vision_vocab_size, |
| 356 | dtype=self.dtype, |
| 357 | param_dtype=self.param_dtype, |
| 358 | use_bias=False, |
| 359 | kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range), |
| 360 | precision=self.precision, |
| 361 | ) |
| 362 | self.lm_head = nn.Dense( |
| 363 | self.config.vocab_size, |
| 364 | dtype=self.dtype, |
| 365 | param_dtype=self.param_dtype, |
| 366 | use_bias=False, |
| 367 | kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range), |
| 368 | precision=self.precision, |
| 369 | ) |
| 370 | |
| 371 | def __call__( |
| 372 | self, |
nothing calls this directly
no test coverage detected