Function test

examples/python_plugin/build_lookup.py:23–57

(dtype)

                               dtype=torch.int32).cuda()

    def test(dtype):
        builder = tensorrt_llm.Builder()
        builder.strongly_typed = True
        network = builder.create_network()
        with tensorrt_llm.net_guard(network):
            x = Tensor(
                name="x",
                shape=index_shape,
                dtype=tensorrt_llm.str_dtype_to_trt("int32"),
            )
            y = Tensor(name="y",
                       shape=(vocab_size, n_embed),
                       dtype=torch_dtype_to_trt(dtype))

            def lookup(x, y):
                lookup_plugin = LookUpPlugin(False, True)
                return lookup_plugin(x, y)

            output = lookup(x, y)

            output.mark_output("output", torch_dtype_to_str(torch.float32))

        builder_config = builder.create_builder_config("float32")
        engine = builder.build_engine(network, builder_config)
        assert engine is not None

        output_dir = Path("tmp") / torch_dtype_to_str(dtype)
        output_dir.mkdir(parents=True, exist_ok=True)

        engine_path = output_dir / "lookup.engine"
        config_path = output_dir / "config.json"

        with engine_path.open("wb") as f:
            f.write(engine)
        builder.save_config(builder_config, str(config_path))

    test(torch.bfloat16)
    test(torch.float16)