fix: patch pypi whl version labels (#916)

.github/workflows/pypi-gpu.yml

@@ -36,6 +36,8 @@ jobs:
 
       - name: Set pyproject.toml version to match github tag
         shell: bash
+        env:
+          RELEASE_TAG: ${{ github.ref_name }}
         run: |
             mv pyproject.toml pyproject.toml.orig
             sed "s/ezkl/ezkl-gpu/" pyproject.toml.orig >pyproject.toml

.github/workflows/pypi.yml

@@ -32,6 +32,14 @@ jobs:
           python-version: 3.12
           architecture: x64
 
+      - name: Set pyproject.toml version to match github tag
+        shell: bash
+        env:
+          RELEASE_TAG: ${{ github.ref_name }}
+        run: |
+          mv pyproject.toml pyproject.toml.orig
+          sed "s/0\\.0\\.0/${RELEASE_TAG//v}/" pyproject.toml.orig >pyproject.toml
+
       - name: Set Cargo.toml version to match github tag
         shell: bash
         env:
@@ -89,6 +97,14 @@ jobs:
           python-version: 3.12
           architecture: ${{ matrix.target }}
 
+      - name: Set pyproject.toml version to match github tag
+        shell: bash
+        env:
+          RELEASE_TAG: ${{ github.ref_name }}
+        run: |
+          mv pyproject.toml pyproject.toml.orig
+          sed "s/0\\.0\\.0/${RELEASE_TAG//v}/" pyproject.toml.orig >pyproject.toml
+
       - name: Set Cargo.toml version to match github tag
         shell: bash
         env:
@@ -138,6 +154,14 @@ jobs:
           python-version: 3.12
           architecture: x64
 
+      - name: Set pyproject.toml version to match github tag
+        shell: bash
+        env:
+          RELEASE_TAG: ${{ github.ref_name }}
+        run: |
+          mv pyproject.toml pyproject.toml.orig
+          sed "s/0\\.0\\.0/${RELEASE_TAG//v}/" pyproject.toml.orig >pyproject.toml
+
       - name: Set Cargo.toml version to match github tag
         shell: bash
         env:
@@ -148,7 +172,6 @@ jobs:
           mv Cargo.lock Cargo.lock.orig
           sed "s/0\\.0\\.0/${RELEASE_TAG//v}/" Cargo.lock.orig >Cargo.lock
 
-
       - name: Install required libraries
         shell: bash
         run: |
@@ -187,57 +210,6 @@ jobs:
           name: wheels
           path: dist
 
-  #  There's a problem with the maturin-action toolchain for arm arch leading to failed builds
-  # linux-cross:
-  #   runs-on: ubuntu-latest
-  #   strategy:
-  #     matrix:
-  #       target: [aarch64, armv7]
-  #   steps:
-  #   - uses: actions/checkout@v4
-  #   - uses: actions/setup-python@v4
-  #     with:
-  #       python-version: 3.12
-
-  #   - name: Install cross-compilation tools for aarch64
-  #     if: matrix.target == 'aarch64'
-  #     run: |
-  #       sudo apt-get update
-  #       sudo apt-get install -y gcc make gcc-aarch64-linux-gnu binutils-aarch64-linux-gnu libc6-dev-arm64-cross libusb-1.0-0-dev libatomic1-arm64-cross
-
-  #   - name: Install cross-compilation tools for armv7
-  #     if: matrix.target == 'armv7'
-  #     run: |
-  #       sudo apt-get update
-  #       sudo apt-get install -y gcc make gcc-arm-linux-gnueabihf binutils-arm-linux-gnueabihf libc6-dev-armhf-cross libusb-1.0-0-dev libatomic1-armhf-cross
-
-  #   - name: Build wheels
-  #     uses: PyO3/maturin-action@v1
-  #     with:
-  #       target: ${{ matrix.target }}
-  #       manylinux: auto
-  #       args: --release --out dist --features python-bindings
-
-  #   - uses: uraimo/run-on-arch-action@v2.5.0
-  #     name: Install built wheel
-  #     with:
-  #       arch: ${{ matrix.target }}
-  #       distro: ubuntu20.04
-  #       githubToken: ${{ github.token }}
-  #       install: |
-  #         apt-get update
-  #         apt-get install -y --no-install-recommends python3 python3-pip
-  #         pip3 install -U pip
-  #       run: |
-  #         pip3 install ezkl --no-index --find-links dist/ --force-reinstall
-  #         python3 -c "import ezkl"
-
-  #   - name: Upload wheels
-  #     uses: actions/upload-artifact@v3
-  #     with:
-  #       name: wheels
-  #       path: dist
-
   musllinux:
     permissions:
       contents: read
@@ -273,6 +245,7 @@ jobs:
           sed "s/0\\.0\\.0/${RELEASE_TAG//v}/" Cargo.toml.orig >Cargo.toml
           mv Cargo.lock Cargo.lock.orig
           sed "s/0\\.0\\.0/${RELEASE_TAG//v}/" Cargo.lock.orig >Cargo.lock
+
       - name: Install required libraries
         shell: bash
         run: |
@@ -323,6 +296,14 @@ jobs:
         with:
           python-version: 3.12
 
+      - name: Set pyproject.toml version to match github tag
+        shell: bash
+        env:
+          RELEASE_TAG: ${{ github.ref_name }}
+        run: |
+          mv pyproject.toml pyproject.toml.orig
+          sed "s/0\\.0\\.0/${RELEASE_TAG//v}/" pyproject.toml.orig >pyproject.toml
+
       - name: Set Cargo.toml version to match github tag
         shell: bash
         env:
@@ -366,8 +347,6 @@ jobs:
     permissions:
       id-token: write
     if: "startsWith(github.ref, 'refs/tags/')"
-    # TODO: Uncomment if linux-cross is working
-    # needs: [ macos, windows, linux, linux-cross, musllinux, musllinux-cross ]
     needs: [macos, windows, linux, musllinux, musllinux-cross]
     steps:
       - uses: actions/download-artifact@v3
@@ -375,24 +354,20 @@ jobs:
           name: wheels
       - name: List Files
         run: ls -R
-
-      # Both publish steps will fail if there is no trusted publisher setup
-      # On failure the publish step will then simply continue to the next one
+      
+      # # publishes to TestPyPI
+      # - name: Publish package distribution to TestPyPI
+      #   uses: pypa/gh-action-pypi-publish@unstable/v1
+      #   with:
+      #     repository-url: https://test.pypi.org/legacy/
+      #     packages-dir: ./
 
       # publishes to PyPI
       - name: Publish package distributions to PyPI
-        continue-on-error: true
         uses: pypa/gh-action-pypi-publish@unstable/v1
         with:
           packages-dir: ./
 
-      # publishes to TestPyPI
-      - name: Publish package distribution to TestPyPI
-        continue-on-error: true
-        uses: pypa/gh-action-pypi-publish@unstable/v1
-        with:
-          repository-url: https://test.pypi.org/legacy/
-          packages-dir: ./
 
   doc-publish:
     permissions:
@@ -409,4 +384,4 @@ jobs:
         with:
           webhook_url: ${{ secrets.RTDS_WEBHOOK_URL }}
           webhook_token: ${{ secrets.RTDS_WEBHOOK_TOKEN }}
-          commit_ref: ${{ github.ref_name }}
+          commit_ref: ${{ github.ref_name }}

src/circuit/ops/chip.rs

@@ -592,9 +592,9 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> BaseConfig<F> {
                         // this is 0 if the index is the same as the column index (starting from 1)
 
                         let col_expr = sel.clone()
-                            * table
+                            * (table
                                 .selector_constructor
-                                .get_expr_at_idx(col_idx, synthetic_sel);
+                                .get_expr_at_idx(col_idx, synthetic_sel));
 
                         let multiplier =
                             table.selector_constructor.get_selector_val_at_idx(col_idx);
@@ -626,6 +626,40 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> BaseConfig<F> {
                         res
                     });
                 }
+
+                // add a degree-k custom constraint of the following form to the range check and
+                // static lookup configuration.
+                // 𝑚𝑢𝑙𝑡𝑖𝑠𝑒𝑙 · ∏ (𝑠𝑒𝑙 − 𝑖) = 0 where 𝑠𝑒𝑙 is the synthetic_sel, and the product is over the set of overflowed columns
+                // and 𝑚𝑢𝑙𝑡𝑖𝑠𝑒𝑙 is the selector value at the column index
+                cs.create_gate("range_check_on_sel", |cs| {
+                    let synthetic_sel = match len {
+                        1 => Expression::Constant(F::from(1)),
+                        _ => match index {
+                            VarTensor::Advice { inner: advices, .. } => {
+                                cs.query_advice(advices[x][y], Rotation(0))
+                            }
+                            _ => unreachable!(),
+                        },
+                    };
+
+                    let range_check_on_synthetic_sel = match len {
+                        1 => Expression::Constant(F::from(0)),
+                        _ => {
+                            let mut initial_expr = Expression::Constant(F::from(1));
+                            for i in 0..len {
+                                initial_expr = initial_expr
+                                    * (synthetic_sel.clone()
+                                        - Expression::Constant(F::from(i as u64)))
+                            }
+                            initial_expr
+                        }
+                    };
+
+                    let sel = cs.query_selector(multi_col_selector);
+
+                    Constraints::with_selector(sel, vec![range_check_on_synthetic_sel])
+                });
+
                 self.static_lookups
                     .selectors
                     .insert((nl.clone(), x, y), multi_col_selector);
@@ -904,9 +938,9 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> BaseConfig<F> {
                         let default_x = range_check.get_first_element(col_idx);
 
                         let col_expr = sel.clone()
-                            * range_check
+                            * (range_check
                                 .selector_constructor
-                                .get_expr_at_idx(col_idx, synthetic_sel);
+                                .get_expr_at_idx(col_idx, synthetic_sel));
 
                         let multiplier = range_check
                             .selector_constructor
@@ -929,6 +963,40 @@ impl<F: PrimeField + TensorType + PartialOrd + std::hash::Hash> BaseConfig<F> {
                         res
                     });
                 }
+
+                // add a degree-k custom constraint of the following form to the range check and
+                // static lookup configuration.
+                // 𝑚𝑢𝑙𝑡𝑖𝑠𝑒𝑙 · ∏ (𝑠𝑒𝑙 − 𝑖) = 0 where 𝑠𝑒𝑙 is the synthetic_sel, and the product is over the set of overflowed columns
+                // and 𝑚𝑢𝑙𝑡𝑖𝑠𝑒𝑙 is the selector value at the column index
+                cs.create_gate("range_check_on_sel", |cs| {
+                    let synthetic_sel = match len {
+                        1 => Expression::Constant(F::from(1)),
+                        _ => match index {
+                            VarTensor::Advice { inner: advices, .. } => {
+                                cs.query_advice(advices[x][y], Rotation(0))
+                            }
+                            _ => unreachable!(),
+                        },
+                    };
+
+                    let range_check_on_synthetic_sel = match len {
+                        1 => Expression::Constant(F::from(0)),
+                        _ => {
+                            let mut initial_expr = Expression::Constant(F::from(1));
+                            for i in 0..len {
+                                initial_expr = initial_expr
+                                    * (synthetic_sel.clone()
+                                        - Expression::Constant(F::from(i as u64)))
+                            }
+                            initial_expr
+                        }
+                    };
+
+                    let sel = cs.query_selector(multi_col_selector);
+
+                    Constraints::with_selector(sel, vec![range_check_on_synthetic_sel])
+                });
+
                 self.range_checks
                     .selectors
                     .insert((range, x, y), multi_col_selector);

src/circuit/ops/layouts.rs

@@ -75,7 +75,7 @@ fn optimum_convex_function<F: PrimeField + TensorType + PartialOrd + std::hash::
     region: &mut RegionCtx<F>,
     x: &ValTensor<F>,
     f: impl Fn(&BaseConfig<F>, &mut RegionCtx<F>, &ValTensor<F>) -> Result<ValTensor<F>, CircuitError>,
-) -> Result<(), CircuitError> {
+) -> Result<ValTensor<F>, CircuitError> {
     let one = create_constant_tensor(F::from(1), 1);
 
     let f_x = f(config, region, x)?;
@@ -87,22 +87,17 @@ fn optimum_convex_function<F: PrimeField + TensorType + PartialOrd + std::hash::
     let f_x_minus_1 = f(config, region, &x_minus_1)?;
 
     // because the function is convex, the result should be the minimum of the three
-    // not that we offset the x by 1 to get the next value
-    // f(x) <= f(x+1) and f(x) <= f(x-1)
+    // note that we offset the x by 1 to get the next value
+    // f(x) <= f(x+1) and f(x) < f(x-1)
     // the result is 1 if the function is optimal solely because of the convexity of the function
-    // the distances can be equal but this is only possible if f(x) and f(x+1) are both optimal (or f(x) and f(x-1)).
+    // the distances can be equal but this is only possible if f(x) and f(x+1) are both optimal, but if (f(x) = f(x + 1))
+    // f(x+1) is not smaller than f(x + 1 - 1) = f(x) and thus f(x) is unique
     let f_x_is_opt_rhs = less_equal(config, region, &[f_x.clone(), f_x_plus_1.clone()])?;
-    let f_x_is_opt_lhs = less_equal(config, region, &[f_x.clone(), f_x_minus_1.clone()])?;
+    let f_x_is_opt_lhs = less(config, region, &[f_x.clone(), f_x_minus_1.clone()])?;
 
     let is_opt = and(config, region, &[f_x_is_opt_lhs, f_x_is_opt_rhs])?;
 
-    let mut comparison_unit = create_constant_tensor(F::ONE, is_opt.len());
-    comparison_unit.reshape(is_opt.dims())?;
-
-    // assert that the result is 1
-    enforce_equality(config, region, &[is_opt, comparison_unit])?;
-
-    Ok(())
+    Ok(is_opt)
 }
 
 /// Err is less than some constant
@@ -290,7 +285,14 @@ pub(crate) fn recip<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
         Ok(distance)
     };
 
-    optimum_convex_function(config, region, &claimed_output, err_func)?;
+    // we need to add 1 to the points where it is zero to ignore the cvx opt conditions at those points
+    let mut is_opt = optimum_convex_function(config, region, &claimed_output, err_func)?;
+    is_opt = pairwise(config, region, &[is_opt, equal_zero_mask], BaseOp::Add)?;
+
+    let mut comparison_unit = create_constant_tensor(F::ONE, is_opt.len());
+    comparison_unit.reshape(is_opt.dims())?;
+    // assert that the result is 1
+    enforce_equality(config, region, &[is_opt, comparison_unit])?;
 
     Ok(claimed_output)
 }
@@ -362,7 +364,13 @@ pub fn sqrt<F: PrimeField + TensorType + PartialOrd + std::hash::Hash>(
         Ok(distance)
     };
 
-    optimum_convex_function(config, region, &claimed_output, err_func)?;
+    let is_opt = optimum_convex_function(config, region, &claimed_output, err_func)?;
+
+    let mut comparison_unit = create_constant_tensor(F::ONE, is_opt.len());
+    comparison_unit.reshape(is_opt.dims())?;
+
+    // assert that the result is 1
+    enforce_equality(config, region, &[is_opt, comparison_unit])?;
 
     Ok(claimed_output)
 }