//  [python]
def call(
        self,
        image_embeddings: tf.Tensor,
        image_positional_embeddings: tf.Tensor,
        sparse_prompt_embeddings: tf.Tensor,
        dense_prompt_embeddings: tf.Tensor,
        multimask_output: bool,
        output_attentions: Optional[bool] = None,
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        batch_size, num_channels, height, width = shape_list(image_embeddings)
        point_batch_size = tf.math.maximum(1, tf.shape(sparse_prompt_embeddings)[1])

        output_tokens = tf.concat([self.iou_token, self.mask_tokens], axis=0)  # Should be (1, 32) + (4, 32) = (5, 32)
        output_tokens = tf.tile(
            output_tokens[None, None, :], [batch_size, point_batch_size, 1, 1]
        )  # Should be (batch_size, point_size, 5, 32)

        # Matt: The original Torch code checked that the sum of sparse_prompt_embeddings equalled 0. However, this only
        #       happens when the sparse prompt embeddings are an empty tensor with shape[1] == 0. I replaced
        #       it with an explicit shape check to avoid data-dependent control flow which breaks XLA.
        if shape_list(sparse_prompt_embeddings)[1] != 0:
            tokens = tf.concat((output_tokens, sparse_prompt_embeddings), axis=2)
        else:
            tokens = output_tokens
        point_embeddings = tf.cast(tokens, self.iou_token.dtype)

        image_embeddings = image_embeddings + dense_prompt_embeddings
        image_embeddings = tf.repeat(image_embeddings, point_batch_size, axis=0)
        image_positional_embeddings = tf.repeat(image_positional_embeddings, point_batch_size, axis=0)

        point_embedding, image_embeddings, attentions = self.transformer(
            point_embeddings=point_embeddings,
            image_embeddings=image_embeddings,
            image_positional_embeddings=image_positional_embeddings,
            output_attentions=output_attentions,
        )
        iou_token_out = point_embedding[:, :, 0, :]
        mask_tokens_out = point_embedding[:, :, 1 : (1 + self.num_mask_tokens), :]

        image_embeddings = tf.transpose(image_embeddings, perm=(0, 1, 3, 2))
        image_embeddings = tf.reshape(image_embeddings, [batch_size * point_batch_size, num_channels, height, width])

        upscaled_embedding = self.upscale_conv1(image_embeddings)
        upscaled_embedding = self.activation(self.upscale_layer_norm(upscaled_embedding))
        upscaled_embedding = self.activation(self.upscale_conv2(upscaled_embedding))

        hyper_in_list = []
        for i in range(self.num_mask_tokens):
            current_mlp = self.output_hypernetworks_mlps[i]
            hyper_in_list += [current_mlp(mask_tokens_out[:, :, i, :])]
        hyper_in = tf.stack(hyper_in_list, axis=2)

        _, num_channels, height, width = shape_list(upscaled_embedding)
        upscaled_embedding = tf.reshape(
            upscaled_embedding, [batch_size, point_batch_size, num_channels, height * width]
        )
        masks = tf.reshape(hyper_in @ upscaled_embedding, [batch_size, point_batch_size, -1, height, width])

        iou_pred = self.iou_prediction_head(iou_token_out)

        if multimask_output:
            mask_slice = slice(1, None)
        else:
            mask_slice = slice(0, 1)
        masks = masks[:, :, mask_slice, :, :]
        iou_pred = iou_pred[:, :, mask_slice]

        outputs = (masks, iou_pred)

        if output_attentions:
            outputs = outputs + (attentions,)
        else:
            outputs = outputs + (None,)

        return outputs