_merge_input_ids_with_image_features({
        inputs_embeds,
        image_features,
        input_ids,
        attention_mask,
    }) {

        // @ts-expect-error TS2339
        const image_token_index = this.config.image_token_index;

        const idsList = input_ids.tolist();

        // NOTE: we use .findIndex instead of .indexOf to perform weak comparison (==) between BigInt and Number
        const indexOfImage = idsList.map(x => x.findIndex(x => x == image_token_index));

        const noImages = indexOfImage.every(x => x === -1);
        const allImages = indexOfImage.every(x => x !== -1);
        if (!noImages && !allImages) {
            // Check for padding reasons
            throw new Error('Every input should contain either 0 or 1 image token.');
        }

        if (noImages) {
            return {
                inputs_embeds,
                attention_mask,
            }
        }

        const stacked = [];
        const stacked_attention_mask = [];
        for (let i = 0; i < indexOfImage.length; ++i) {
            const index = indexOfImage[i];

            const e = inputs_embeds[i];
            const im = image_features[i];
            const am = attention_mask[i];
            stacked.push(
                cat([
                    e.slice([0, index]),
                    im,
                    e.slice([index + 1, e.dims[0]]),
                ], 0)
            );

            stacked_attention_mask.push(
                cat([
                    am.slice([0, index]),
                    ones([im.dims[0]]),
                    am.slice([index + 1, am.dims[0]])
                ], 0)
            )
        }

        return {
            inputs_embeds: stack(stacked, 0),
            attention_mask: stack(stacked_attention_mask, 0),
        }
    }
The LLAVA model which consists of a vision backbone and a language model.