Merge pull request #25538 from jakevdp:raise-runtimeerror

PiperOrigin-RevId: 707200618
jax-ml · Dec 17, 2024 · 4aebe55 · 4aebe55
2 parents 0fa5419 + 2518c62
commit 4aebe55
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/_src/core.py b/jax/_src/core.py
@@ -1669,7 +1669,7 @@ def get_sharding(sharding, ndim):
 
   context_mesh = mesh_lib.get_abstract_mesh()
   if not context_mesh:
-    return RuntimeError("Please set the mesh via `jax.set_mesh` API.")
+    raise RuntimeError("Please set the mesh via `jax.set_mesh` API.")
   assert sharding is None
   return NamedSharding(context_mesh, P(*[None] * ndim))