Add jax lowering to cast f64 to f32 for aten ops add, sub, lt, mul, where.

PiperOrigin-RevId: 790985881

Author: junjiang-lab

ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py

@@ -16,6 +16,7 @@
 from typing import Any, Callable
 from ai_edge_torch import fx_infra
 from ai_edge_torch import lowertools
+from ai_edge_torch.odml_torch import optimization_barrier as optimization_barrier_lib
 import torch
 import torch.utils._pytree as pytree
 
@@ -276,6 +277,7 @@ def embedding(*args, **kwargs):
     # Explicitly reshape back to the original shape. This places the ReshapeOp
     # outside of the HLFB.
     output = torch.reshape(output, (*(original_idx_shape), embedding_dim))
+    output, _ = optimization_barrier_lib.optimization_barrier(output, idx)
     return output
 
   node.target = embedding

ai_edge_torch/_convert/test/test_convert.py

@@ -576,6 +576,24 @@ def forward(self, x: torch.Tensor):
       self.fail(f"Conversion failed with bloat16 inputs: {err}")
     # pylint: enable=broad-except
 
+  def test_convert_model_with_torch_linspace_operation(self):
+    """Test converting a simple model with torch.linspace operation."""
+
+    class SampleModel(nn.Module):
+
+      def forward(self, x: torch.Tensor):
+        return torch.linspace(0.5, 10.5, steps=x.shape[0], dtype=torch.float64)
+
+    model = SampleModel().eval()
+    args = (torch.randint(0, 100, (10, 10), dtype=torch.int64),)
+
+    try:
+      # Expect this to fix the error during conversion
+      ai_edge_torch.convert(model, args)
+    except Exception as err:
+      self.fail(f"Conversion failed with int64 inputs: {err}")
+    # pylint: enable=broad-except
+
   def test_compile_model(self):
     """Tests AOT compilation of a simple Add module."""
 

ai_edge_torch/generative/examples/gemma/gemma1.py

@@ -23,7 +23,7 @@
 import torch
 from torch import nn
 
-TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+TENSOR_NAMES_FUSED_QKV = loading_utils.ModelLoader.TensorNames(
     ff_up_proj="model.layers.{}.mlp.up_proj",
     ff_down_proj="model.layers.{}.mlp.down_proj",
     ff_gate_proj="model.layers.{}.mlp.gate_proj",
@@ -36,6 +36,24 @@
     lm_head=None,
 )
 
+TENSOR_NAMES_SEP_QKV = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="model.layers.{}.mlp.up_proj",
+    ff_down_proj="model.layers.{}.mlp.down_proj",
+    ff_gate_proj="model.layers.{}.mlp.gate_proj",
+    attn_query_proj="model.layers.{}.self_attn.q_proj",
+    attn_key_proj="model.layers.{}.self_attn.k_proj",
+    attn_value_proj="model.layers.{}.self_attn.v_proj",
+    attn_output_proj="model.layers.{}.self_attn.o_proj",
+    pre_attn_norm="model.layers.{}.input_layernorm",
+    post_attn_norm="model.layers.{}.post_attention_layernorm",
+    embedding="model.embed_tokens",
+    final_norm="model.norm",
+)
+
+TENSOR_NAMES_DICT = {
+    "safetensors": TENSOR_NAMES_SEP_QKV,
+    "kaggle": TENSOR_NAMES_FUSED_QKV,
+}
 
 class Gemma1(model_builder.DecoderOnlyModel):
   """A Gemma1 model built from the Edge Generative API layers."""
@@ -94,11 +112,28 @@ def build_2b_model(
     custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
     mask_cache_size: int = 0,
 ) -> nn.Module:
-  return model_builder.build_decoder_only_model(
-      checkpoint_path=checkpoint_path,
-      config=get_model_config_2b(),
-      tensor_names=TENSOR_NAMES,
-      model_class=Gemma1,
-      custom_loader=custom_loader,
-      mask_cache_size=mask_cache_size,
+
+  # A list to store the reasons for each failure
+  key_errors = []
+
+  for tensor_names in TENSOR_NAMES_DICT.values():
+    try:
+      return model_builder.build_decoder_only_model(
+          checkpoint_path=checkpoint_path,
+          config=get_model_config_2b(),
+          tensor_names=tensor_names,
+          model_class=Gemma1,
+          custom_loader=custom_loader,
+          mask_cache_size=mask_cache_size,
+      )
+    except KeyError as ke:
+      # Store the specific key that was missing for later
+      key_errors.append(f"Missing key: {ke}")
+      continue
+
+  # If the loop finishes, raise an error with all the collected details
+  error_details = "\n".join(key_errors)
+  raise RuntimeError(
+      "Failed to build model after trying all configurations. "
+      f"Encountered the following errors:\n{error_details}"
   )

ai_edge_torch/generative/examples/smolvlm2/verify_encoder.py

@@ -0,0 +1,95 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Verifies the reauthored SmolVLM2 Image Encoder model."""
+
+import logging
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.smolvlm2 import smolvlm2
+from ai_edge_torch.generative.examples.smolvlm2 import vision_encoder
+from PIL import Image
+import requests
+import torch
+import transformers
+
+_IMAGE_URL = flags.DEFINE_string(
+    "image_url",
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true",
+    "The image URI to encode.",
+)
+
+_CHECKPOINT = flags.DEFINE_string(
+    "checkpoint",
+    "HuggingFaceTB/SmolVLM2-2.2B-Instruct",
+    "The checkpoint to verify.",
+)
+
+_REAUTHORTHED_CHECKPOINT = flags.DEFINE_string(
+    "pretrained_weights",
+    None,
+    "The path to the model's pretrained weights.",
+)
+
+
+def main(_):
+  checkpoint = _CHECKPOINT.value
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_model = transformers.AutoModelForImageTextToText.from_pretrained(
+      checkpoint
+  )
+  original_model = original_model.eval().model
+
+  logging.info("Building the reauthored checkpoint from: %s", checkpoint)
+  reauthored_checkpoint = _REAUTHORTHED_CHECKPOINT.value
+  if reauthored_checkpoint is None:
+    raise ValueError("reauthored_checkpoint is required.")
+
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = vision_encoder.build_image_encoder(reauthored_checkpoint)
+
+  logging.info("Loading the tokenizer from: %s", checkpoint)
+  processor = transformers.AutoProcessor.from_pretrained(checkpoint)
+
+  logging.info("Loading the image from: %s", _IMAGE_URL.value)
+  image = Image.open(requests.get(_IMAGE_URL.value, stream=True).raw)
+  pixel_values = processor(images=image, return_tensors="pt")["pixel_values"]
+
+  logging.info("Forwarding the original model...")
+  outputs_original = original_model.get_image_features(pixel_values)
+  logging.info("outputs_original's shape: %s", outputs_original.shape)
+
+  pixel_values = pixel_values.reshape(
+      pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:]
+  )
+  logging.info("Forwarding the reauthored model...")
+  outputs_reauthored = reauthored_model.forward(
+      pixel_values=pixel_values
+  )
+  logging.info("outputs_reauthored's shape: %s", outputs_reauthored.shape)
+
+  try:
+    assert torch.allclose(
+        outputs_original, outputs_reauthored, atol=1e-03, rtol=1e-04
+    )
+  except AssertionError as e:
+    logging.error("*** FAILED *** verify with an image")
+    raise e
+  else:
+    logging.info("*** PASSED *** verify with an image")
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/smolvlm2/vision_encoder.py

@@ -20,7 +20,7 @@
 """
 
 from dataclasses import dataclass
-from typing import Callable, Dict
+from typing import Callable, Dict, Optional
 
 from ai_edge_torch.generative.examples.paligemma import image_encoder
 import ai_edge_torch.generative.layers.model_config as cfg
@@ -127,9 +127,20 @@ def __init__(
   def forward(
       self,
       pixel_values: torch.Tensor,
-      export_config: export_cfg.ExportConfig = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> torch.Tensor:
-    x = self.siglip_encoder(pixel_values)
+    # Embed the image according to SiplipVisionEmbeddings.
+    x = self.siglip_encoder.tok_embedding(pixel_values)
+    x = x.flatten(2).transpose(1, 2)
+    x = x + self.siglip_encoder.tok_embedding_position
+
+    # Pass a dummy mask because SDPA attention impl expects non-None mask.
+    mask = torch.zeros(x.shape[0], 1, x.shape[1], x.shape[1])
+    for _, block in enumerate(self.siglip_encoder.transformer_blocks):
+      x = block(x, mask=mask)
+    x = self.siglip_encoder.final_norm(x)
+
+    # Project the image embeddings to text hidden size.
     x = self.connector(x)
     return x
 
@@ -166,7 +177,7 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       output_proj_use_bias=True,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.LAYER_NORM, epsilon=1e-6
+      type=cfg.NormalizationType.LAYER_NORM, epsilon=1e-6,
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.SEQUENTIAL,
@@ -189,15 +200,13 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       image_embedding=image_embedding_config,
       block_configs=block_config,
       final_norm_config=norm_config,
-      # num_mm_tokens_per_image=81,
-      # enable_hlfb=False
   )
   return config
 
 
 def build_image_encoder(
     checkpoint_path: str,
-    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    custom_loader: Optional[Callable[[str], Dict[str, torch.Tensor]]] = None,
 ) -> FullVisionEncoder:
   """Builds a FullVisionEncoder from the checkpoint path."""
   encoder_config = get_image_encoder_config()
@@ -208,7 +217,6 @@ def build_image_encoder(
   )
   loader.load(encoder.siglip_encoder, strict=False)
 
-  loader = loading_utils.ModelLoader(checkpoint_path, None, custom_loader)
   state = loader.get_state()
   converted_state = dict()
   converted_state["modality_projection.weight"] = state.pop(

ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py

@@ -180,6 +180,11 @@ def _aten_rsqrt_decomp(x):
   return torch.ops.tfl.rsqrt(x)
 
 
+@register_decomp(torch.ops.aten.neg.default)
+def _aten_neg_decomp(x):
+  return torch.ops.tfl.neg(x)
+
+
 @register_decomp(torch.ops.aten.gelu.default)
 def _aten_gelu_decomp(x, approximate="none"):
   return torch.ops.tfl.gelu(x, approximate != "none")
@@ -317,6 +322,38 @@ def _aten_select_int_decomp(x, dim, index):
   return torch.ops.tfl.squeeze(sliced, [dim])
 
 
+@register_decomp(torch.ops.aten.slice.Tensor)
+def _aten_slice_tensor_decomp(x, dim=0, start=None, end=None, step=1):
+  rank = x.dim()
+  dim_size = x.shape[dim]
+
+  # Initialize begin, end, strides for tfl.strided_slice
+  begin = [0] * rank
+  end_vec = list(x.shape)
+  strides = [1] * rank
+
+  # The logic below is to match PyTorch's `slice` behavior.
+  # `start` and `end` can be negative, which means they count from the end.
+  # `start=None` defaults to 0.
+  # `end=None` or a large number defaults to `dim_size` after clamping.
+
+  start_val = 0 if start is None else start
+  if start_val < 0:
+    start_val += dim_size
+
+  end_val = dim_size if end is None else end
+  if end_val < 0:
+    end_val += dim_size
+
+  # Clamp start and end to be within the dimension size, following PyTorch's
+  # logic.
+  start_val = max(0, min(start_val, dim_size))
+  end_val = max(start_val, min(end_val, dim_size))
+
+  begin[dim], end_vec[dim], strides[dim] = start_val, end_val, step
+  return torch.ops.tfl.strided_slice(x, begin, end_vec, strides)
+
+
 @register_decomp(torch.ops.aten.where.self)
 def _aten_where_self_decomp(condition, x, y):
   x, y = _promote_types_for_binary_op(x, y)
@@ -351,3 +388,27 @@ def _aten__softmax_decomp(
     softmax_result = torch.ops.tfl.softmax(x_permuted)
     # Transpose the result back to the original dimensions.
     return torch.ops.tfl.transpose(softmax_result, dims)
+
+
+@register_decomp(torch.ops.aten.topk.default)
+def _aten_topk_decomp(self, k, dim=-1, largest=True, sorted=True):
+  if not largest:
+    raise ValueError("Only largest=True is supported for torch.topk.")
+
+  if dim < 0:
+    dim = self.dim() + dim
+
+  if dim != self.dim() - 1:
+    self = torch.transpose(self, dim, -1)
+
+  # Ignores sorted value: tfl.topk_v2 only supports sorted=True, but it doesn't
+  # affect the correctness of the output.
+  out, indices = torch.ops.tfl.topk_v2(self, k)
+
+  if dim != self.dim() - 1:
+    out = torch.transpose(out, dim, -1)
+    indices = torch.transpose(indices, dim, -1)
+
+  # torch.topk returns int64 indices, but tfl.topk_v2 returns indices in int32.
+  indices = indices.to(torch.int64)
+  return out, indices

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -286,6 +286,18 @@ def _tfl_rsqrt_lowering(
   )
 
 
+@lower(torch.ops.tfl.neg.default)
+def _tfl_neg_lowering(
+    lctx: LoweringContext,
+    x: ir.Value,
+) -> ir.Value:
+  return _ir_operation(
+      "tfl.neg",
+      results=lowering_utils.node_meta_to_ir_types(lctx.node),
+      operands=[x],
+  )
+
+
 @lower(torch.ops.tfl.gelu.default)
 def _tfl_gelu_lowering(
     lctx: LoweringContext,
@@ -674,3 +686,20 @@ def _tfl_softmax_lowering(
           "beta": ir.FloatAttr.get(ir.F32Type.get(), beta),
       },
   )
+
+
+@lower(torch.ops.tfl.topk_v2.default)
+def _tfl_topk_v2_lowering(
+    lctx: LoweringContext,
+    x: ir.Value,
+    k: int,
+) -> tuple[ir.Value, ir.Value]:
+  return _ir_operation(
+      "tfl.topk_v2",
+      results=lowering_utils.node_meta_to_ir_types(lctx.node),
+      operands=[
+          x,
+          lowering_utils.numpy_array_constant(np.array(k, dtype=np.int32)),
+      ],
+      attributes={},
+  )