Merge branch 'main' into f8E4M3_f8E3M4

stablehlo/integrations/python/tests/stablehlo.py

@@ -301,7 +301,6 @@ def test_reference_api():
 def test_get_smaller_version():
   curr_version = stablehlo.get_current_version()
   min_version = stablehlo.get_minimum_version()
-  print(curr_version)
   assert stablehlo.get_smaller_version(curr_version, min_version) == min_version
 
 

stablehlo/testdata/quantized/convolution_qi8.mlir

@@ -0,0 +1,17 @@
+// RUN: stablehlo-translate --interpret -split-input-file %s
+
+
+module attributes {jax.uses_shape_polymorphism = true} {
+  func.func @main() -> tensor<i1> {
+    %cst = stablehlo.constant dense<[[[[1.000000e+00], [2.000000e+00], [5.000000e+00], [6.000000e+00]], [[3.000000e+00], [4.000000e+00], [7.000000e+00], [8.000000e+00]], [[1.000000e+01], [1.100000e+01], [1.400000e+01], [1.500000e+01]], [[1.200000e+01], [1.300000e+01], [1.600000e+01], [1.700000e+01]]]]> : tensor<1x4x4x1xf32>
+    %cst_0 = stablehlo.constant dense<1.000000e+00> : tensor<3x3x1x1xf32>
+    %cst_1 = stablehlo.constant dense<2.47494364> : tensor<1x2x2x1xf32>
+    %0 = stablehlo.uniform_quantize %cst_0 : (tensor<3x3x1x1xf32>) -> tensor<3x3x1x1x!quant.uniform<i8:f32, 0.0039188104517319626:-128>>
+    %1 = stablehlo.uniform_quantize %cst : (tensor<1x4x4x1xf32>) -> tensor<1x4x4x1x!quant.uniform<i8:f32, 0.0039172410964965817:-128>>
+    %2 = stablehlo.convolution(%1, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [4, 4], lhs_dilate = [2, 2]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x4x4x1x!quant.uniform<i8:f32, 0.0039172410964965817:-128>>, tensor<3x3x1x1x!quant.uniform<i8:f32, 0.0039188104517319626:-128>>) -> tensor<1x2x2x1x!quant.uniform<i32:f32, 1.5350925350904778E-5>>
+    %3 = stablehlo.uniform_quantize %2 : (tensor<1x2x2x1x!quant.uniform<i32:f32, 1.5350925350904778E-5>>) -> tensor<1x2x2x1x!quant.uniform<i8:f32, 0.0097056613248937273:-128>>
+    %4 = stablehlo.uniform_dequantize %3 : (tensor<1x2x2x1x!quant.uniform<i8:f32, 0.0097056613248937273:-128>>) -> tensor<1x2x2x1xf32>
+    %5 = stablehlo.custom_call @check.eq(%cst_1, %4) : (tensor<1x2x2x1xf32>, tensor<1x2x2x1xf32>) -> tensor<i1>
+    return %5 : tensor<i1>
+  }
+}

stablehlo/testdata/quantized/dot_general_dtypes_and_precision_lhs_float32_1_3_4__rhs_float32_1_4_3__dimensionnumbers____2_1___1_2-3746161505671304111_qi8.mlir

@@ -0,0 +1,16 @@
+// RUN: stablehlo-translate --interpret -split-input-file %s
+
+module attributes {jax.uses_shape_polymorphism = true} {
+  func.func @main() -> tensor<i1> {
+    %cst = stablehlo.constant dense<[[[2.59234142, -2.35737705, 3.07461166, -5.77705336], [-3.64460349, 0.689637601, -0.0876942202, 1.62593222], [-4.65739489, 0.247004092, -1.08101177, -0.238710642]]]> : tensor<1x3x4xf32>
+    %cst_0 = stablehlo.constant dense<[[[-3.64053726, 1.4618907, -0.867068588], [2.86438012, -2.70172548, -0.4580172], [0.140140817, -0.666462898, -7.101100e-01], [1.02142382, 0.236523077, 0.760420739]]]> : tensor<1x4x3xf32>
+    %cst_1 = stablehlo.constant dense<-0.558793128> : tensor<1xf32>
+    %0 = stablehlo.uniform_quantize %cst_0 : (tensor<1x4x3xf32>) -> tensor<1x4x3x!quant.uniform<i8:f32, 0.0039153145808799592>>
+    %1 = stablehlo.uniform_quantize %cst : (tensor<1x3x4xf32>) -> tensor<1x3x4x!quant.uniform<i8:f32, 0.0039166809297075458>>
+    %2 = stablehlo.dot_general %1, %0, batching_dims = [0] x [0], contracting_dims = [2, 1] x [1, 2], precision = [HIGHEST, HIGHEST] : (tensor<1x3x4x!quant.uniform<i8:f32, 0.0039166809297075458>>, tensor<1x4x3x!quant.uniform<i8:f32, 0.0039153145808799592>>) -> tensor<1x!quant.uniform<i32:f32, 1.533503795273843E-5>>
+    %3 = stablehlo.uniform_quantize %2 : (tensor<1x!quant.uniform<i32:f32, 1.533503795273843E-5>>) -> tensor<1x!quant.uniform<i8:f32, 0.0051740104076909085:-20>>
+    %4 = stablehlo.uniform_dequantize %3 : (tensor<1x!quant.uniform<i8:f32, 0.0051740104076909085:-20>>) -> tensor<1xf32>
+    %5 = stablehlo.custom_call @check.eq(%cst_1, %4) : (tensor<1xf32>, tensor<1xf32>) -> tensor<i1>
+    return %5 : tensor<i1>
+  }
+}

stablehlo/testdata/quantized/dot_general_dtypes_and_precision_lhs_float32_1_3_4__rhs_float32_1_4_3__dimensionnumbers____2_1___1_2-601699960869302464_qi8.mlir

@@ -0,0 +1,16 @@
+// RUN: stablehlo-translate --interpret -split-input-file %s
+
+module attributes {jax.uses_shape_polymorphism = true} {
+  func.func @main() -> tensor<i1> {
+    %cst = stablehlo.constant dense<[[[3.21390557, 2.41580057, -0.537137687, -1.02739561], [1.35577726, -2.48765302, 3.99296689, 3.09424305], [-2.41700459, -1.63692343, -6.27982473, 2.19841671]]]> : tensor<1x3x4xf32>
+    %cst_0 = stablehlo.constant dense<[[[-0.227901727, -0.527938426, 2.04744601], [0.46251452, -2.02832699, -5.24830675], [-1.4312098, -5.60030842, 6.4256258], [-0.754353106, -7.84103918, -1.26745594]]]> : tensor<1x4x3xf32>
+    %cst_1 = stablehlo.constant dense<-0.372528762> : tensor<1xf32>
+    %0 = stablehlo.uniform_quantize %cst_0 : (tensor<1x4x3xf32>) -> tensor<1x4x3x!quant.uniform<i8:f32, 0.0039153145808799592>>
+    %1 = stablehlo.uniform_quantize %cst : (tensor<1x3x4xf32>) -> tensor<1x3x4x!quant.uniform<i8:f32, 0.0039166809297075458>>
+    %2 = stablehlo.dot_general %1, %0, batching_dims = [0] x [0], contracting_dims = [2, 1] x [1, 2] : (tensor<1x3x4x!quant.uniform<i8:f32, 0.0039166809297075458>>, tensor<1x4x3x!quant.uniform<i8:f32, 0.0039153145808799592>>) -> tensor<1x!quant.uniform<i32:f32, 1.533503795273843E-5>>
+    %3 = stablehlo.uniform_quantize %2 : (tensor<1x!quant.uniform<i32:f32, 1.533503795273843E-5>>) -> tensor<1x!quant.uniform<i8:f32, 0.0051740104076909085:-20>>
+    %4 = stablehlo.uniform_dequantize %3 : (tensor<1x!quant.uniform<i8:f32, 0.0051740104076909085:-20>>) -> tensor<1xf32>
+    %5 = stablehlo.custom_call @check.eq(%cst_1, %4) : (tensor<1xf32>, tensor<1xf32>) -> tensor<i1>
+    return %5 : tensor<i1>
+  }
+}

stablehlo/testdata/quantized/dot_general_dtypes_and_precision_lhs_float32_1_3_4__rhs_float32_1_4_3__dimensionnumbers____2_1___1_26892231390636518787_qi8.mlir

@@ -0,0 +1,16 @@
+// RUN: stablehlo-translate --interpret -split-input-file %s
+
+module attributes {jax.uses_shape_polymorphism = true} {
+  func.func @main() -> tensor<i1> {
+    %cst = stablehlo.constant dense<[[[3.29589868, 1.23571205, 1.77729869, -1.18627894], [-0.334112585, -0.701673209, 1.37458372, 1.33988047], [-1.23167276, -1.47567272, -5.089730e+00, 4.67841911]]]> : tensor<1x3x4xf32>
+    %cst_0 = stablehlo.constant dense<[[[-1.98014688, 6.42148352, -4.60602903], [0.598645091, -1.56104231, 1.48972511], [2.90995979, -0.593875766, -4.25129318], [-1.8781029, 1.97643542, 2.91169143]]]> : tensor<1x4x3xf32>
+    %cst_1 = stablehlo.constant dense<0.760579526> : tensor<1xf32>
+    %0 = stablehlo.uniform_quantize %cst_0 : (tensor<1x4x3xf32>) -> tensor<1x4x3x!quant.uniform<i8:f32, 0.0039153145808799592>>
+    %1 = stablehlo.uniform_quantize %cst : (tensor<1x3x4xf32>) -> tensor<1x3x4x!quant.uniform<i8:f32, 0.0039166809297075458>>
+    %2 = stablehlo.dot_general %1, %0, batching_dims = [0] x [0], contracting_dims = [2, 1] x [1, 2], precision = [HIGH, HIGH] : (tensor<1x3x4x!quant.uniform<i8:f32, 0.0039166809297075458>>, tensor<1x4x3x!quant.uniform<i8:f32, 0.0039153145808799592>>) -> tensor<1x!quant.uniform<i32:f32, 1.533503795273843E-5>>
+    %3 = stablehlo.uniform_quantize %2 : (tensor<1x!quant.uniform<i32:f32, 1.533503795273843E-5>>) -> tensor<1x!quant.uniform<i8:f32, 0.0051740104076909085:-20>>
+    %4 = stablehlo.uniform_dequantize %3 : (tensor<1x!quant.uniform<i8:f32, 0.0051740104076909085:-20>>) -> tensor<1xf32>
+    %5 = stablehlo.custom_call @check.eq(%cst_1, %4) : (tensor<1xf32>, tensor<1xf32>) -> tensor<i1>
+    return %5 : tensor<i1>
+  }
+}

stablehlo/testdata/quantized/dot_general_dtypes_and_precision_lhs_float32_1_3_4__rhs_float32_1_4_3__dimensionnumbers____2_1___1_2_____0____0_qi8.mlir

@@ -0,0 +1,16 @@
+// RUN: stablehlo-translate --interpret -split-input-file %s
+
+module attributes {jax.uses_shape_polymorphism = true} {
+  func.func @main() -> tensor<i1> {
+    %cst = stablehlo.constant dense<[[[-1.02716792, -0.84180355, -1.21495497, 1.22473526], [1.533764, 3.72483683, -3.63489795, 1.59403718], [-2.40816736, -1.31785822, -1.78571689, 0.202816486]]]> : tensor<1x3x4xf32>
+    %cst_0 = stablehlo.constant dense<[[[2.14761162, 2.50920868, -4.64897203], [1.72184336, -1.84444284, -1.14315307], [-1.28549385, 2.20468378, 1.76219583], [-3.49526167, -2.27656484, 2.42410755]]]> : tensor<1x4x3xf32>
+    %cst_1 = stablehlo.constant dense<-0.558793128> : tensor<1xf32>
+    %0 = stablehlo.uniform_quantize %cst_0 : (tensor<1x4x3xf32>) -> tensor<1x4x3x!quant.uniform<i8:f32, 0.0039153145808799592>>
+    %1 = stablehlo.uniform_quantize %cst : (tensor<1x3x4xf32>) -> tensor<1x3x4x!quant.uniform<i8:f32, 0.0039166809297075458>>
+    %2 = stablehlo.dot_general %1, %0, batching_dims = [0] x [0], contracting_dims = [2, 1] x [1, 2] : (tensor<1x3x4x!quant.uniform<i8:f32, 0.0039166809297075458>>, tensor<1x4x3x!quant.uniform<i8:f32, 0.0039153145808799592>>) -> tensor<1x!quant.uniform<i32:f32, 1.533503795273843E-5>>
+    %3 = stablehlo.uniform_quantize %2 : (tensor<1x!quant.uniform<i32:f32, 1.533503795273843E-5>>) -> tensor<1x!quant.uniform<i8:f32, 0.0051740104076909085:-20>>
+    %4 = stablehlo.uniform_dequantize %3 : (tensor<1x!quant.uniform<i8:f32, 0.0051740104076909085:-20>>) -> tensor<1xf32>
+    %5 = stablehlo.custom_call @check.eq(%cst_1, %4) : (tensor<1xf32>, tensor<1xf32>) -> tensor<i1>
+    return %5 : tensor<i1>
+  }
+}

stablehlo/testdata/quantized/round_nearest_afz_qi8.mlir

@@ -0,0 +1,13 @@
+// RUN: stablehlo-translate --interpret -split-input-file %s
+
+module attributes {jax.uses_shape_polymorphism = true} {
+  func.func @main() -> tensor<i1> {
+    %cst = stablehlo.constant dense<[-2.500000e+00, 4.000000e-01, 5.000000e-01, 6.000000e-01, 2.500000e+00]> : tensor<5xf32>
+    %cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<5xf32>
+    %0 = stablehlo.uniform_quantize %cst : (tensor<5xf32>) -> tensor<5x!quant.uniform<i8:f32, 0.0038905945478701124:-1>>
+    %1 = stablehlo.round_nearest_afz %0 : (tensor<5x!quant.uniform<i8:f32, 0.0038905945478701124:-1>>) -> tensor<5x!quant.uniform<i8:f32, 7.843137254901961E-9>>
+    %2 = stablehlo.uniform_dequantize %1 : (tensor<5x!quant.uniform<i8:f32, 7.843137254901961E-9>>) -> tensor<5xf32>
+    %3 = stablehlo.custom_call @check.eq(%cst_0, %2) : (tensor<5xf32>, tensor<5xf32>) -> tensor<i1>
+    return %3 : tensor<i1>
+  }
+}

stablehlo/testdata/quantized/round_rounding_methods_shape_float32_2_5__roundingmethod_0_qi8.mlir

@@ -0,0 +1,13 @@
+// RUN: stablehlo-translate --interpret -split-input-file %s
+
+module attributes {jax.uses_shape_polymorphism = true} {
+  func.func @main() -> tensor<i1> {
+    %cst = stablehlo.constant dense<[[5.000000e-01, 1.200000e+00, 1.500000e+00, 1.700000e+00, 2.500000e+00], [-5.000000e-01, -1.200000e+00, -1.500000e+00, -1.700000e+00, -2.500000e+00]]> : tensor<2x5xf32>
+    %cst_0 = stablehlo.constant dense<[[0.996078491, 0.996078491, 0.996078491, 0.996078491, 0.996078491], [-0.996078491, -0.996078491, -0.996078491, -0.996078491, -0.996078491]]> : tensor<2x5xf32>
+    %0 = stablehlo.uniform_quantize %cst : (tensor<2x5xf32>) -> tensor<2x5x!quant.uniform<i8:f32, 0.0078306291617599184>>
+    %1 = stablehlo.round_nearest_afz %0 : (tensor<2x5x!quant.uniform<i8:f32, 0.0078306291617599184>>) -> tensor<2x5x!quant.uniform<i8:f32, 0.0078431372549019607:-1>>
+    %2 = stablehlo.uniform_dequantize %1 : (tensor<2x5x!quant.uniform<i8:f32, 0.0078431372549019607:-1>>) -> tensor<2x5xf32>
+    %3 = stablehlo.custom_call @check.eq(%cst_0, %2) : (tensor<2x5xf32>, tensor<2x5xf32>) -> tensor<i1>
+    return %3 : tensor<i1>
+  }
+}