Rename with_dll_constraint to with_layout_constraint for the upcoming Layout API rename!

PiperOrigin-RevId: 753636449

Author: yashk2810

jax/_src/pjit.py

@@ -3120,49 +3120,49 @@ def use_explicit_axes(*axes):
   with mesh_lib.use_abstract_mesh(new_mesh):
     yield
 
-# -------------------- with_dll_constraint --------------------
+# -------------------- with_layout_constraint --------------------
 
-def with_dll_constraint(x, layouts):
+def with_layout_constraint(x, layouts):
   x_flat, tree = tree_flatten(x)
-  layouts_flat = tuple(flatten_axes("with_dll_constraint layouts", tree,
+  layouts_flat = tuple(flatten_axes("with_layout_constraint layouts", tree,
                                     layouts))
   if any(not isinstance(l, DeviceLocalLayout) for l in layouts_flat):
     raise ValueError(
-        'layouts passed to `with_dll_constraint` must be of type'
+        'layouts passed to `with_layout_constraint` must be of type'
         f' `DeviceLocalLayout`. Got {[type(l) for l in layouts_flat]}')
   check_aval_layout_compatibility(
       layouts_flat, x_flat, ("",) * len(layouts_flat),
-      "with_dll_constraint arguments")
-  outs = [dll_constraint_p.bind(xf, layout=l)
+      "with_layout_constraint arguments")
+  outs = [layout_constraint_p.bind(xf, layout=l)
           for xf, l in zip(x_flat, layouts_flat)]
   return tree_unflatten(tree, outs)
 
-dll_constraint_p = core.Primitive('dll_constraint')
-dll_constraint_p.def_abstract_eval(lambda x, **_: x)
-ad.deflinear2(dll_constraint_p,
-              lambda ct, _, **params: (dll_constraint_p.bind(ct, **params),))
+layout_constraint_p = core.Primitive('layout_constraint')
+layout_constraint_p.def_abstract_eval(lambda x, **_: x)
+ad.deflinear2(layout_constraint_p,
+              lambda ct, _, **params: (layout_constraint_p.bind(ct, **params),))
 
-def _dll_constraint_impl(x, *, layout):
+def _layout_constraint_impl(x, *, layout):
   if not isinstance(x, xc.ArrayImpl):
     raise ValueError(
-        'with_dll_constraint in eager mode can only be applied to'
+        'with_layout_constraint in eager mode can only be applied to'
         f' jax.Arrays. Got {type(x)}')
   if x.layout.device_local_layout == layout:  # type: ignore
     return x
   return api.jit(_identity_fn, out_shardings=Layout(layout, x.sharding))(x)
-dll_constraint_p.def_impl(_dll_constraint_impl)
+layout_constraint_p.def_impl(_layout_constraint_impl)
 
-def _dll_constraint_hlo_lowering(ctx, x_node, *, layout):
+def _layout_constraint_hlo_lowering(ctx, x_node, *, layout):
   aval, = ctx.avals_in
   out_aval, = ctx.avals_out
   return [mlir.wrap_with_layout_op(ctx, x_node, out_aval, layout, aval)]
-mlir.register_lowering(dll_constraint_p,
-                       _dll_constraint_hlo_lowering)
+mlir.register_lowering(layout_constraint_p,
+                       _layout_constraint_hlo_lowering)
 
-def _dll_constraint_batcher(axis_data, vals_in, dims_in, layout):
+def _layout_constraint_batcher(axis_data, vals_in, dims_in, layout):
   raise NotImplementedError
-batching.fancy_primitive_batchers[dll_constraint_p] = _dll_constraint_batcher
-batching.skippable_batchers[dll_constraint_p] = lambda _: ()
+batching.fancy_primitive_batchers[layout_constraint_p] = _layout_constraint_batcher
+batching.skippable_batchers[layout_constraint_p] = lambda _: ()
 
 # -------------------- helpers --------------------
 

jax/experimental/jax2tf/tests/primitives_test.py

@@ -178,7 +178,7 @@ def test_primitive_coverage(self):
         continue
       if p.name == "sharding_constraint":
         continue
-      if p.name == "dll_constraint":
+      if p.name == "layout_constraint":
         continue
       if p.name == "mesh_cast":
         continue

jax/experimental/layout.py

@@ -17,5 +17,5 @@
     Layout as Layout,
 )
 from jax._src.pjit import (
-    with_dll_constraint as with_dll_constraint,
+    with_layout_constraint as with_layout_constraint,
 )

tests/layout_test.py

@@ -23,7 +23,7 @@
 from jax._src import config
 from jax._src import test_util as jtu
 from jax._src.util import safe_zip
-from jax.experimental.layout import (with_dll_constraint, Layout,
+from jax.experimental.layout import (with_layout_constraint, Layout,
                                      DeviceLocalLayout as DLL)
 from jax.experimental.compute_on import compute_on
 
@@ -745,7 +745,7 @@ def f(x):
     self.assertArraysEqual(out, np_inp * 2)
     self.assertEqual(out.layout, out_layout)
 
-  def test_with_dll_constraint(self):
+  def test_with_layout_constraint(self):
     if not jtu.test_device_matches(['tpu']):
       self.skipTest('Only works for TPU')
     mesh = jtu.create_mesh((2, 2), ('x', 'y'))
@@ -761,7 +761,7 @@ def f(x):
       y = x.T
       # Constrain `y` to the original layout of `arr` because without it,
       # the layout of `y` would be the transpose of `arr`.
-      y = with_dll_constraint(y, custom_dll)
+      y = with_layout_constraint(y, custom_dll)
       return y * 2
 
     f(arr)  # doesn't crash