Method __call__

src/models/text_decoder.py:351–411

(self, x, u, deterministic=True)

    @nn.compact
    def __call__(self, x, u, deterministic=True):
        out = {}
        dpr = [
            float(x) for x in np.linspace(
                0,
                self.drop_path,
                self.depth)]  # drop path decay
        # Input Encoder
        CrossAttnBlockLayer = CrossAttnEncoder1DBlock
        if self.remat_policy not in (None, "none"):
            logging.info(f"remat policy: {self.remat_policy}")
            if self.remat_policy == "minimal":
                policy = jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims
            else:
                policy = None
            logging.info(f"activation checkpointing {self.remat_policy}")
            CrossAttnBlockLayer = remat(  # pylint: disable=invalid-name
                CrossAttnEncoder1DBlock, prevent_cse=True, policy=policy, static_argnums=(3,)
            )  # "deterministic" is a static argument in CrossAttnEncoder1DBlock

        BlockLayer = Encoder1DBlock
        if self.remat_policy not in (None, "none"):
            logging.info(f"remat policy: {self.remat_policy}")
            if self.remat_policy == "minimal":
                policy = jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims
            else:
                policy = None
            logging.info(f"activation checkpointing {self.remat_policy}")
            BlockLayer = remat(  # pylint: disable=invalid-name
                Encoder1DBlock, prevent_cse=True, policy=policy, static_argnums=(1,)
            )  # "deterministic" is a static argument in Encoder1DBlock

        for lyr in range(self.depth):
            x, out[f"block{lyr:02d}"] = BlockLayer(
                name=f"encoderblock_{lyr}",
                mlp_dim=self.mlp_dim,
                depth=self.depth,
                num_heads=self.num_heads,
                dropout=self.dropout,
                drop_path=dpr[lyr],
                casual_mask=self.casual_mask,
                use_flash_attn=self.use_flash_attn,
                dtype=self.dtype,
                param_dtype=self.param_dtype,
                mesh=self.mesh
            )(x, deterministic)
            x, out[f"crossattn_block{lyr:02d}"] = CrossAttnBlockLayer(
                name=f"crossattn_encoderblock_{lyr}",
                mlp_dim=self.mlp_dim, 
                depth=self.depth,
                num_heads=self.num_heads, 
                dropout=self.dropout,
                drop_path=dpr[lyr],
                casual_mask=False,
                use_flash_attn=self.use_flash_attn,
                dtype=self.dtype,
                param_dtype=self.param_dtype,
                mesh=self.mesh)(x, u, None, deterministic)