Enforce JAX_RELEASE version

2 files changed, 7 insertions, 3 deletions

diff --git a/.SRCINFO b/.SRCINFO
index 4b8a27a991d5..abc4764f9f8d 100644
--- a/.SRCINFO
+++ b/.SRCINFO
@@ -1,7 +1,7 @@
 pkgbase = python-jax
 	pkgdesc = Differentiate, compile, and transform Numpy code.
 	pkgver = 0.4.17
-	pkgrel = 1
+	pkgrel = 2
 	url = https://github.com/google/jax/
 	arch = any
 	groups = jax
diff --git a/PKGBUILD b/PKGBUILD
index 9f434ea0cd47..04e319b67310 100644
--- a/PKGBUILD
+++ b/PKGBUILD
@@ -2,7 +2,7 @@
 
 pkgname='python-jax'
 pkgver=0.4.17
-pkgrel=1
+pkgrel=2
 pkgdesc='Differentiate, compile, and transform Numpy code.'
 arch=('any')
 url='https://github.com/google/jax/'
@@ -19,6 +19,10 @@ makedepends=('python-build' 'python-installer' 'python-setuptools' 'python-wheel
 source=("jax-${pkgver}.tar.gz::https://github.com/google/jax/archive/refs/tags/jax-v${pkgver}.tar.gz")
 sha256sums=('31d8e55b6c2655f57d4f5214e4620834d484fcdd783e33f51ec83fa2a4700918')
 
+prepare() {
+    export JAX_RELEASE=$pkgver
+}
+
 build() {
     cd jax-jax-v$pkgver
     python -m build -n -w
@@ -28,5 +32,5 @@ package() {
     python -m installer \
         --compile-bytecode 1 \
         --destdir $pkgdir \
-        jax-jax-v$pkgver/dist/jax-*.whl
+        jax-jax-v$pkgver/dist/jax-$pkgver-*.whl
 }