Update references to JAX's GitHub repo

JAX has moved from https://github.com/google/jax to https://github.com/jax-ml/jax PiperOrigin-RevId: 702886997
google · Dec 10, 2024 · d0f65b2 · d0f65b2
1 parent 2f606dc
commit d0f65b2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -1,7 +1,7 @@
 # Jaxite
 
 Jaxite is a fully homomorphic encryption backend targeting TPUs and GPUs,
-written in [JAX](https://github.com/google/jax).
+written in [JAX](https://github.com/jax-ml/jax).
 
 It implements the [CGGI cryptosystem](https://eprint.iacr.org/2018/421) with
 some optimizations, and is a supported backend for [Google's FHE

diff --git a/jaxite/jaxite_lib/matrix_utils.py b/jaxite/jaxite_lib/matrix_utils.py
@@ -181,7 +181,7 @@ def hpmatmul_conv_adapt_conv(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
 
   lhs: jax.Array = jax.lax.bitcast_convert_type(x, new_dtype=jnp.uint8)  # bnmp
   rhs: jax.Array = jax.lax.bitcast_convert_type(y, new_dtype=jnp.uint8)  # nk1q
-  # https://github.com/google/jax/issues/11483
+  # https://github.com/jax-ml/jax/issues/11483
   rhs = jax.lax.rev(rhs, [2])
   # rhs = jlax.rev(rhs, dimensions=[3])