Promote operand types before calling TFLite binary ops.

PiperOrigin-RevId: 773888968

Author: junjiang-lab

ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py

@@ -66,33 +66,60 @@ def _aten_sub_tensor_decomp(x, y, alpha=1):
   return out
 
 
+def _promote_types_for_binary_op(x, y):
+  """Promotes operand types for a binary op."""
+  # TFLite's binary ops require operands to have the same element type.
+  # We promote the types before calling the op.
+  if not isinstance(x, torch.Tensor):
+    # x is a scalar, y must be a tensor.
+    x = torch.scalar_tensor(x, dtype=y.dtype)
+  elif not isinstance(y, torch.Tensor):
+    # y is a scalar, x is a tensor.
+    # Handle scalar operand by converting scalar to a tensor.
+    # The dtype of the new tensor should match x's dtype for promotion.
+    y = torch.scalar_tensor(y, dtype=x.dtype)
+
+  target_dtype = torch.promote_types(x.dtype, y.dtype)
+  if x.dtype != target_dtype:
+    x = x.to(target_dtype)
+  if y.dtype != target_dtype:
+    y = y.to(target_dtype)
+  return x, y
+
+
 @register_decomp(torch.ops.aten.mul.Tensor)
 def _aten_mul_tensor_decomp(x, y):
+  x, y = _promote_types_for_binary_op(x, y)
   return torch.ops.tfl.mul(x, y)
 
 
 @register_decomp(torch.ops.aten.mul.Scalar)
 def _aten_mul_scalar_decomp(x, y):
+  x, y = _promote_types_for_binary_op(x, y)
   return torch.ops.tfl.mul(x, y)
 
 
 @register_decomp(torch.ops.aten.div.Tensor)
 def _aten_div_tensor_decomp(x, y):
+  x, y = _promote_types_for_binary_op(x, y)
   return torch.ops.tfl.div(x, y)
 
 
 @register_decomp(torch.ops.aten.pow.Scalar)
 def _aten_pow_scalar_decomp(x, y):
+  x, y = _promote_types_for_binary_op(x, y)
   return torch.ops.tfl.pow(x, y)
 
 
 @register_decomp(torch.ops.aten.pow.Tensor_Scalar)
 def _aten_pow_tensor_scalar_decomp(x, y):
+  x, y = _promote_types_for_binary_op(x, y)
   return torch.ops.tfl.pow(x, y)
 
 
 @register_decomp(torch.ops.aten.pow.Tensor_Tensor)
 def _aten_pow_tensor_tensor_decomp(x, y):
+  x, y = _promote_types_for_binary_op(x, y)
   return torch.ops.tfl.pow(x, y)
 
 
@@ -117,11 +144,13 @@ def _aten_mean_dim_decomp(x, dim, keepdim=False):
 
 @register_decomp(torch.ops.aten.gt.Tensor)
 def _aten_gt_tensor_decomp(x, y):
+  x, y = _promote_types_for_binary_op(x, y)
   return torch.ops.tfl.greater(x, y)
 
 
 @register_decomp(torch.ops.aten.lt.Tensor)
 def _aten_lt_tensor_decomp(x, y):
+  x, y = _promote_types_for_binary_op(x, y)
   return torch.ops.tfl.less(x, y)
 
 
@@ -280,6 +309,7 @@ def _aten_select_int_decomp(x, dim, index):
 
 @register_decomp(torch.ops.aten.where.self)
 def _aten_where_self_decomp(condition, x, y):
+  x, y = _promote_types_for_binary_op(x, y)
   return torch.ops.tfl.select_v2(condition, x, y)
 
 

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -429,22 +429,32 @@ def _tfl_range_lowering(
       lowering_utils.torch_dtype_to_ir_element_type(output_torch_dtype)
   )
 
+  # All operands and the output of tfl.range must have the same element type.
+  # We cast all operands to the expected output element type.
   operands = []
-  for val_py_scalar in [
-      start,
-      limit,
-      delta,
-  ]:
-    if isinstance(val_py_scalar, ir.Value):
-      operands.append(val_py_scalar)
-    else:
+  for operand in [start, limit, delta]:
+    if not isinstance(operand, ir.Value):
+      # Convert python scalars to ir.Value.
       numpy_scalar_0d = (
-          torch.tensor(val_py_scalar, dtype=output_torch_dtype)
-          .detach()
-          .numpy()
+          torch.tensor(operand, dtype=output_torch_dtype).detach().numpy()
       )
-      scalar_tensor_val = lowering_utils.numpy_array_constant(numpy_scalar_0d)
-      operands.append(scalar_tensor_val)
+      operand = lowering_utils.numpy_array_constant(numpy_scalar_0d)
+
+    # `operand` is now an ir.Value.
+    # Cast its element type to the output element type if they don't match.
+    operand_type = operand.type
+    if not isinstance(operand_type, ir.RankedTensorType):
+      raise TypeError(
+          "tfl.range operand expected to be RankedTensorType, got"
+          f" {operand_type}"
+      )
+
+    if operand_type.element_type != tflite_op_internal_element_type:
+      cast_to_type = ir.RankedTensorType.get(
+          operand_type.shape, tflite_op_internal_element_type
+      )
+      operand = stablehlo.convert(cast_to_type, operand)
+    operands.append(operand)
 
   # Define the result type that the tfl.range *kernel* (the custom op) will
   # produce.

ai_edge_torch/odml_torch/experimental/torch_tfl/test/test_torch_tfl_impls.py

@@ -132,6 +132,7 @@ def _assert_export_and_close(
       ("aten_div_Tensor_2", torch.ops.aten.div.Tensor, (rnd(torch.float32, (10, 10)), np.random.rand(),), dict()),
       ("aten_pow_Scalar_0", torch.ops.aten.pow.Scalar, (np.random.rand(), rnd(torch.float32, (10, 10)),), dict()),
       ("aten_pow_Tensor_Scalar_0", torch.ops.aten.pow.Tensor_Scalar, (rnd(torch.float32, (10, 10)), np.random.rand(),), dict()),
+      ("aten_pow_Tensor_Scalar_1", torch.ops.aten.pow.Tensor_Scalar, (rnd(torch.float32, (10, 10)), np.random.randint(2, 5),), dict()),
       ("aten_pow_Tensor_Tensor_0", torch.ops.aten.pow.Tensor_Tensor, (rnd(torch.float32, (10, 10)), rnd(torch.float32, (10, 10)),), dict()),
       ("aten_bitwise_and_Tensor_0", torch.ops.aten.bitwise_and.Tensor, (rnd(torch.bool, (10, 10)), rnd(torch.bool, (10, 10)),), dict()),
       ("aten_mean_dim_0", torch.ops.aten.mean.dim, (rnd(torch.float32, (10, 10)), 0), dict()),