(self, image_embeds, text_embeds, *, train=False)
| 435 | |
| 436 | @nn.compact |
| 437 | def __call__(self, image_embeds, text_embeds, *, train=False): |
| 438 | out = {} |
| 439 | |
| 440 | if self.drop_token > 0: |
| 441 | image_embeds = image_embeds[:, :image_embeds.shape[1] - self.drop_token + 1] |
| 442 | |
| 443 | ni, li, di = image_embeds.shape |
| 444 | nt, lt, dt = text_embeds.shape |
| 445 | assert ni == nt, f'the image embed is {image_embeds.shape} and text embed is {text_embeds.shape}' |
| 446 | |
| 447 | image_projection_layer = nn.Dense( |
| 448 | self.width, |
| 449 | name="image_projection_layer", |
| 450 | use_bias=False, |
| 451 | kernel_init=nn.initializers.normal(stddev=di ** -0.5)) |
| 452 | image_embeds = out["projected_image_embeds"] = image_projection_layer(image_embeds) |
| 453 | |
| 454 | text_projection_layer = nn.Dense( |
| 455 | self.width, |
| 456 | name="text_projection_layer", |
| 457 | use_bias=False, |
| 458 | kernel_init=nn.initializers.normal(stddev=dt ** -0.5)) |
| 459 | text_embeds = out["projected_text_embeds"] = text_projection_layer(text_embeds) |
| 460 | |
| 461 | learnable_tokens = self.param( |
| 462 | 'learnable_tokens', |
| 463 | nn.initializers.normal(stddev=1.0), |
| 464 | (self.num_learnable_tokens, self.width) |
| 465 | ) |
| 466 | learnable_tokens = jnp.tile(learnable_tokens[None, :, :], (ni, 1, 1)) |
| 467 | |
| 468 | image_embeds = nn.with_logical_constraint(image_embeds, ("activation_batch", "activation_length", "activation_embed")) |
| 469 | text_embeds = nn.with_logical_constraint(text_embeds, ("activation_batch", "activation_length", "activation_embed")) |
| 470 | |
| 471 | # concatenate image_embeds and learnable_tokens in token dimension |
| 472 | image_embeds = jnp.concatenate([image_embeds, text_embeds], axis=1) |
| 473 | li = image_embeds.shape[1] # update image_embeds token dimension |
| 474 | |
| 475 | # keep text_embeds unchanged or continue processing |
| 476 | # if you need to further process the text embeds, you can add logic here |
| 477 | # currently, text_embeds will remain unchanged |
| 478 | text_embeds = learnable_tokens |
| 479 | lt = text_embeds.shape[1] |
| 480 | |
| 481 | # TODO: figure out if need to do triu |
| 482 | if self.fusion_style == 'concat': |
| 483 | |
| 484 | image_text_embeds = jnp.concatenate((image_embeds, text_embeds), axis=1) |
| 485 | image_text_embeds = nn.with_logical_constraint(image_text_embeds, |
| 486 | ("activation_batch", "activation_length", "activation_embed")) |
| 487 | |
| 488 | # it is actually decoder, but we are importing text_transformer implementation to save effort |
| 489 | decoder_blocks = Encoder( |
| 490 | depth=self.depth, |
| 491 | mlp_dim=self.mlp_dim, |
| 492 | num_heads=self.num_heads, |
| 493 | dropout=self.dropout, |
| 494 | drop_path=self.drop_path, |
nothing calls this directly
no test coverage detected