From 183181afcced85163485f63997bc4fa967daaceb Mon Sep 17 00:00:00 2001
From: Michael Daniels <mdaniels5757@gmail.com>
Date: Fri, 19 Jun 2026 22:06:11 -0400
Subject: [PATCH] python3Packages.flax: fix test with new jax

---
 pkgs/development/python-modules/flax/default.nix | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix
index 89f24260f486..9f653e05715b 100644
--- a/pkgs/development/python-modules/flax/default.nix
+++ b/pkgs/development/python-modules/flax/default.nix
@@ -46,6 +46,14 @@ buildPythonPackage (finalAttrs: {
     hash = "sha256-a78KiTsCCARWZvbxz9QKdUKnjkDJGXcPVVJu5rU4m/U=";
   };
 
+  # DeprecationWarning: `with mesh:` context manager has been deprecated. Please use `with jax.set_mesh(mesh):` instead.
+  postPatch = ''
+    substituteInPlace tests/nnx/transforms_test.py \
+      --replace-fail \
+        "with mesh:" \
+        "with jax.set_mesh(mesh):"
+  '';
+
   build-system = [
     setuptools
     setuptools-scm