ai-edge-torch-nightly 0.6.0.dev20250828__py3-none-any.whl → 0.6.0.dev20250830__py3-none-any.whl

ai_edge_torch/_convert/test/test_convert.py

@@ -594,6 +594,27 @@ class TestConvert(googletest.TestCase):
       self.fail(f"Conversion failed with int64 inputs: {err}")
     # pylint: enable=broad-except
 
+  def test_convert_model_with_torch_div_operation(self):
+    """Test converting a simple model with torch.div operation."""
+
+    class SampleModel(nn.Module):
+
+      def forward(self, x: torch.Tensor, y: torch.Tensor):
+        return x / y
+
+    model = SampleModel().eval()
+    args = (
+        torch.randint(0, 100, (10, 10), dtype=torch.int64),
+        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/odml_torch/lowerings/_jax_lowerings.py

@@ -116,7 +116,6 @@ lower_by_torch_xla2(torch.ops.aten.cosh)
 lower_by_torch_xla2(torch.ops.aten.cumsum)
 lower_by_torch_xla2(torch.ops.aten.detach)
 lower_by_torch_xla2(torch.ops.aten.diagonal)
-lower_by_torch_xla2(torch.ops.aten.div)
 lower_by_torch_xla2(torch.ops.aten.dot)
 lower_by_torch_xla2(torch.ops.aten.embedding)
 lower_by_torch_xla2(torch.ops.aten.empty)
@@ -352,6 +351,8 @@ def _aten_mul_scalar(lctx: LoweringContext, self, other):
     promoted_type = jnp.promote_types(self.dtype, other_dtype)
     if promoted_type == jnp.float64:
       promoted_type = jnp.float32
+    elif promoted_type == jnp.int64:
+      promoted_type = jnp.int32
     return jnp.multiply(
         self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
     )
@@ -369,6 +370,8 @@ def _aten_mul_tensor(lctx: LoweringContext, self, other):
     promoted_type = jnp.promote_types(self.dtype, other_dtype)
     if promoted_type == jnp.float64:
       promoted_type = jnp.float32
+    elif promoted_type == jnp.int64:
+      promoted_type = jnp.int32
     return jnp.multiply(
         self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
     )
@@ -376,6 +379,96 @@ def _aten_mul_tensor(lctx: LoweringContext, self, other):
   return jax_lowering(lctx, self, other)
 
 
+@registry.lower(torch.ops.aten.div.Scalar)
+def _aten_div_scalar(lctx: LoweringContext, self, other):
+  _log_usage(torch.ops.aten.div.Scalar)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    other_dtype = jnp.result_type(other)
+    promoted_type = jnp.promote_types(self.dtype, other_dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    elif promoted_type == jnp.int64:
+      promoted_type = jnp.int32
+    return jnp.divide(
+        self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+    )
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.div.Scalar_mode)
+def _aten_div_scalar_mode(lctx: LoweringContext, self, other, rounding_mode=""):
+  _log_usage(torch.ops.aten.div.Scalar_mode)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    other_dtype = jnp.result_type(other)
+    promoted_type = jnp.promote_types(self.dtype, other_dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    elif promoted_type == jnp.int64:
+      promoted_type = jnp.int32
+    if rounding_mode == "floor":
+      return jnp.floor_divide(
+          self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+      )
+    result = jnp.divide(
+        self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+    )
+    if rounding_mode == "trunc":
+      result = jnp.trunc(result)
+    return result
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.div.Tensor)
+def _aten_div_tensor(lctx: LoweringContext, self, other):
+  _log_usage(torch.ops.aten.div.Tensor)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    other_dtype = jnp.result_type(other)
+    promoted_type = jnp.promote_types(self.dtype, other_dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    elif promoted_type == jnp.int64:
+      promoted_type = jnp.int32
+    return jnp.divide(
+        self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+    )
+
+  return jax_lowering(lctx, self, other)
+
+
+@registry.lower(torch.ops.aten.div.Tensor_mode)
+def _aten_div_tensor_mode(lctx: LoweringContext, self, other, rounding_mode=""):
+  _log_usage(torch.ops.aten.div.Tensor_mode)
+
+  @jax_bridge.wrap
+  def jax_lowering(self, other):
+    other_dtype = jnp.result_type(other)
+    promoted_type = jnp.promote_types(self.dtype, other_dtype)
+    if promoted_type == jnp.float64:
+      promoted_type = jnp.float32
+    elif promoted_type == jnp.int64:
+      promoted_type = jnp.int32
+    if rounding_mode == "floor":
+      return jnp.floor_divide(
+          self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+      )
+    result = jnp.divide(
+        self.astype(promoted_type), jnp.array(other, dtype=promoted_type)
+    )
+    if rounding_mode == "trunc":
+      result = jnp.trunc(result)
+    return result
+
+  return jax_lowering(lctx, self, other)
+
+
 @registry.lower(torch.ops.aten.where.self)
 def _aten_where_self(lctx: LoweringContext, condition, self, other):
   _log_usage(torch.ops.aten.where.self)

ai_edge_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of ai-edge-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.6.0.dev20250828"
+__version__ = "0.6.0.dev20250830"

{ai_edge_torch_nightly-0.6.0.dev20250828.dist-info → ai_edge_torch_nightly-0.6.0.dev20250830.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.6.0.dev20250828
+Version: 0.6.0.dev20250830
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.6.0.dev20250828.dist-info → ai_edge_torch_nightly-0.6.0.dev20250830.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=A7loFu8jE9CsXsfMmHYZ-KDFJiaD8Kkqwm_9d3IVzk0,5638
-ai_edge_torch/version.py,sha256=yl0TcYqmBJlm2Y0OYq0y49hU2ZDXXmQW_OvAUFDKogg,806
+ai_edge_torch/version.py,sha256=hqWeY2tyRT1BA2-SphW3wPQ0c-LtS2ll4xHRheEB2WY,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -27,7 +27,7 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=oXbr9G5Jc21xd1dr2CDrp774I4crs0_kkN490K5fNn0,7312
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/test/test_convert.py,sha256=ioeEpeS07NAs14-nHiPI-6lLTtnALxl8uNtTKvoHdgE,19316
+ai_edge_torch/_convert/test/test_convert.py,sha256=sz2-b3_Tf3tKQgvPqxtu0QR-wgKOuoAKcqz19dUby3g,19962
 ai_edge_torch/_convert/test/test_convert_composites.py,sha256=BCIODgxMI_3MxMLfNWYMGjcz-al-J3z5eDHCiZJXNwY,7992
 ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=6_C2R9--KyNR7_oezZIAfyTSR97tOeEWy4XGcbSxBDE,5778
 ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=1o-gUiwzIuO67FNAJ8DeyKv8fVUeZVNNNwofNVDjYeU,3024
@@ -253,7 +253,7 @@ ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=sC4N5-7RS9yKecs97kM9J56enGvs
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
 ai_edge_torch/odml_torch/lowerings/_decomp_registry.py,sha256=ybOdoFE5HIJTkyiYcc73zpyUyUpioVnAca6k0wyJPs4,2572
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=xUuQjoR0NJhuwG36GuycpKHo9jg783bDSHj9wE4F1Sg,15439
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=Hth8F1SwPVEympR1_T6sEsCmmRXtoO7zatPR1uRLFbY,18411
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
@@ -269,8 +269,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.6.0.dev20250828.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.6.0.dev20250828.dist-info/METADATA,sha256=TiZuiiYOhpawH2NS7795QCs04Vxzoau7dBquFrs0Tkc,2074
-ai_edge_torch_nightly-0.6.0.dev20250828.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.6.0.dev20250828.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.6.0.dev20250828.dist-info/RECORD,,
+ai_edge_torch_nightly-0.6.0.dev20250830.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.6.0.dev20250830.dist-info/METADATA,sha256=ni14Rod4z0CU-YvXBiqbVKbprzGRTC6JlPyH3zQMo78,2074
+ai_edge_torch_nightly-0.6.0.dev20250830.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.6.0.dev20250830.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.6.0.dev20250830.dist-info/RECORD,,