The Google Tensor SDK lets you optimize and run machine learning (ML)
models directly on Google Tensor's dedicated TPU. In addition to standard ML
models, you can compile and accelerate language models within your
LiteRT workflow.

> [!WARNING]
> **Preview:** You can use the `litert-google-tensor` agent skill that automates ahead-of-time (AOT) compilation, on-device benchmarking, Play AI packaging, and Gradle integration directly from natural language prompts in your agent workspace. You can use it for optimizing models on custom Pixel Google Tensor TPUs (Google Tensor G5 / Pixel 10 and Tensor G6 / Pixel 11). See [litert-google-tensor](https://github.com/google-ai-edge/litert-samples/tree/main/skills/litert-google-tensor) for more information.

For certain pre-optimized models, we also provide an extra compiler options
binary file that you can include in your workflow.
To achieve optimal performance on Google Tensor, we recommend using the
following optimization flags during compilation:

## Compilation flags for Google Tensor

Optimize your model compilation process by applying specific
configuration flags to tailor performance and resource usage.
You can use these parameters within your LiteRT Python workflow to adjust
compilation behavior for both PyTorch and TFLite models.

| Flag | Requirement | Description | Default Value |
|---|---|---|---|
| `google_tensor_truncation_type` | Optional | Sets the target data type for floating-point operations. *Supported values:* auto (default), bfloat16, half, no_truncation | auto |
| `google_tensor_sharding_intensity` | Optional | Controls how aggressively the model is split for parallel processing. *Options:* minimal, moderate, extensive, maximum. | minimal |
| `google_tensor_int64_to_int32` | Optional | Set to true to allow the compiler to convert 64-bit integers to 32-bit integers, which may be necessary for some models. | False |
| `google_tensor_enable_large_model_support` | Optional | Set to true if your model is larger than 2GB. | False |
| `google_tensor_enable_4bit_compilation` | Optional | Enables compilation of models with convolution operations containing 4 bit parameters. | False |
| `google_tensor_extra_options_path` | Optional | Passes extra compiler options as a binary file. | "" (Empty) |

You can use these flags with LiteRT Python flow as shown in the following code
snippets:

- When compiling a pytorch model using `ai_edge_torch`

      compiled_models = (
        ai_edge_torch.experimental_add_compilation_backend(tensor_g6_target,
        flag_name1=value1,
        flag_name2=value2, ...).convert(
          channel_last_selfie_segmentation,
          sample_input))

- When compiling a tflite model

      compiled_models = aot_lib.aot_compile(
          tflite_model_path,
          target=[tensor_g6_target],
          flag_name1=value1,
          flag_name2=value2, ...)

## Usage example

In the following example, `google_tensor_truncation_type="half"` flag is used:

      compiled_models = aot_lib.aot_compile(
          tflite_model_path,
          target=[tensor_g6_target],
          keep_going=False,
          google_tensor_truncation_type="half"
      )

See [LiteRT AOT Colab](https://github.com/google-ai-edge/litert-samples/blob/main/samples/litert/colab/LiteRT_AOT_Compilation_Tutorial.ipynb) for more information.

## Compile language models for Google Tensor

To compile language models for Google Tensor, follow the instructions in
[NPU AOT compilation](https://developers.google.com/edge/litert/conversion/pytorch/genai#npu-aot-compilation).

To export LLMs for Google Tensor TPUs, follow the example for the
additional flags required for NPU compilation.

**Example:**

    litert-torch export-hf \
      --model=google/gemma-3-270m-it \
      --output_dir=/tmp/gemma3-270m-google-tensor-g6 \
      --split_cache \
      --externalize_embedder \
      --prefill_lengths=128, \
      --cache_length=1280 \
      --quantization_recipe="weight_only_wi8_afp32" \
      --aot_backend=GOOGLE \
      --aot_soc_model=Tensor_G6 \
      --aot_compilation_config_dict='{"google_tensor_enable_large_model_support": True}'