From 0cfbfe4568d87ff171e6c2fa7691af820511983a Mon Sep 17 00:00:00 2001 From: Albert Torosyan Date: Thu, 3 Jul 2025 11:47:51 +0400 Subject: [PATCH 1/2] [fix] Add check for newer version of jaxlib --- aim/sdk/num_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aim/sdk/num_utils.py b/aim/sdk/num_utils.py index ff16782785..b11b291066 100644 --- a/aim/sdk/num_utils.py +++ b/aim/sdk/num_utils.py @@ -71,6 +71,9 @@ def is_jax_device_array(inst): return True if inst_has_typename(inst, ['jaxlib', 'xla_extension', 'DeviceArray']): return True + if inst_has_typename(inst, ['jaxlib', '_jax', 'ArrayImpl']): + return True + return False From ed0dc4516b5a518943a4c9c099addaffef0d387c Mon Sep 17 00:00:00 2001 From: Albert Torosyan Date: Thu, 3 Jul 2025 11:48:40 +0400 Subject: [PATCH 2/2] [doc] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index bf3ba2d6f6..6fd659e5a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixes: - Fix issues with tag false reassignment (mihran113) +- Fix issue of detecting defice arrays with newer versions of jax (alberttorosyan) ## 3.29.1 May 8, 2025: