From ee6370334543a30c9b603b55b95b58d7ed84881a Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Fri, 28 Aug 2026 12:50:58 -0700 Subject: [PATCH] [JAX] Add an import of jax.experimental.shard_map to modules that use it. Some users were relying on a transitive import via jax.experimental.checkify. jax.experimental.shard_map is itself deprecated but fixing that is for future work. PiperOrigin-RevId: 972741037 --- init2winit/model_lib/dlrm.py | 1 + 1 file changed, 1 insertion(+) diff --git a/init2winit/model_lib/dlrm.py b/init2winit/model_lib/dlrm.py index 888d9f3d..49d4a68b 100644 --- a/init2winit/model_lib/dlrm.py +++ b/init2winit/model_lib/dlrm.py @@ -22,6 +22,7 @@ from init2winit.model_lib import model_utils import jax from jax import nn as jnn +import jax.experimental.shard_map import jax.numpy as jnp from ml_collections.config_dict import config_dict