Rename layout.AUTO to DeviceLocalLayout.AUTO

PiperOrigin-RevId: 621684185

Author: yashk2810

jax/_src/layout.py

@@ -21,9 +21,17 @@
 from jax._src.lib import xla_client as xc
 
 
+class AutoLayout:
+
+  def __repr__(self):
+    return "AUTO"
+
+
 class DeviceLocalLayout:
   layout: xc.PjRtLayout
 
+  AUTO = AutoLayout()
+
   def __init__(self, layout: xc.PjRtLayout):
     self._layout = layout
     self._layout_str = str(self._layout)
@@ -43,14 +51,6 @@ def _to_xla_layout(self) -> str:
     return self._layout_str
 
 
-class AutoLayout:
-
-  def __repr__(self):
-    return "AUTO"
-
-AUTO = AutoLayout()
-
-
 LayoutOptions = Union[DeviceLocalLayout, None, AutoLayout]
 ShardingOptions = Union[Sharding, None, AutoSharding]
 

jax/experimental/layout.py

@@ -14,6 +14,5 @@
 
 from jax._src.layout import (
     DeviceLocalLayout as DeviceLocalLayout,
-    AUTO as AUTO,
     Layout as Layout
 )

tests/layout_test.py

@@ -22,8 +22,7 @@
 import jax.numpy as jnp
 from jax.sharding import NamedSharding, PartitionSpec as P, SingleDeviceSharding
 from jax._src import config
-from jax._src import layout
-from jax._src.layout import Layout
+from jax._src.layout import Layout, DeviceLocalLayout as DLL
 from jax._src import test_util as jtu
 from jax._src.util import safe_zip
 from jax._src import xla_bridge
@@ -90,7 +89,7 @@ def init(x, y):
     sds2 = jax.ShapeDtypeStruct(np_inp2.shape, np_inp2.dtype, sharding=s2)
 
     lowered_apply = jax.jit(apply).lower(
-        sds1, sds2, _in_layouts=layout.AUTO, _out_layouts=layout.AUTO)
+        sds1, sds2, _in_layouts=DLL.AUTO, _out_layouts=DLL.AUTO)
     compiled_apply = lowered_apply.compile()
 
     arg_layouts, kw_layouts = compiled_apply._input_layouts()
@@ -159,8 +158,8 @@ def f(x):
     self.assertArraysEqual(out, np_inp.T)
     self.assertEqual(out.sharding, NamedSharding(mesh, P(None, 'y', 'x')))
 
-    compiled_auto = jax.jit(f).lower(sds, _in_layouts=layout.AUTO,
-                                     _out_layouts=layout.AUTO).compile()
+    compiled_auto = jax.jit(f).lower(sds, _in_layouts=DLL.AUTO,
+                                     _out_layouts=DLL.AUTO).compile()
     self.assertTupleEqual(
         extract_minor_to_major(compiled_auto._input_layouts()[0][0]), (2, 1, 0))
     self.assertTupleEqual(
@@ -177,7 +176,7 @@ def f(x):
       return x.T
 
     compiled = jax.jit(f).lower(
-        arr, _in_layouts=None, _out_layouts=layout.AUTO).compile()
+        arr, _in_layouts=None, _out_layouts=DLL.AUTO).compile()
     self.assertTupleEqual(
         extract_minor_to_major(compiled._input_layouts()[0][0]), (1, 0))
     self.assertTupleEqual(
@@ -195,7 +194,7 @@ def test_sharding_and_layouts(self):
     s = NamedSharding(mesh, P('x', 'y'))
 
     compiled = jax.jit(lambda x: x.T, in_shardings=s, out_shardings=s).lower(
-        np_inp, _in_layouts=layout.AUTO, _out_layouts=layout.AUTO).compile()
+        np_inp, _in_layouts=DLL.AUTO, _out_layouts=DLL.AUTO).compile()
     out = compiled(np_inp)
     self.assertTupleEqual(
         extract_minor_to_major(compiled._input_layouts()[0][0]), (1, 0))
@@ -210,8 +209,8 @@ def f(x, y, z, a, b, c):
 
     shape = (8, 2)
     inps = [np.arange(math.prod(shape)).reshape(shape)] * 6
-    compiled = jax.jit(f).lower(*inps, _in_layouts=layout.AUTO,
-                                _out_layouts=layout.AUTO).compile()
+    compiled = jax.jit(f).lower(*inps, _in_layouts=DLL.AUTO,
+                                _out_layouts=DLL.AUTO).compile()
     arg_layouts, _ = compiled._input_layouts()
     out1, out2 = compiled(*inps)
 
@@ -244,10 +243,10 @@ def f(x):
     with self.assertRaisesRegex(
         ValueError,
         'Layout passed to jit does not match the layout on the respective arg'):
-      jax.jit(f).lower(arr, _in_layouts=layout.AUTO)
+      jax.jit(f).lower(arr, _in_layouts=DLL.AUTO)
 
     compiled = jax.jit(f).lower(
-        sds, _in_layouts=layout.AUTO, _out_layouts=layout.AUTO).compile()
+        sds, _in_layouts=DLL.AUTO, _out_layouts=DLL.AUTO).compile()
 
     with self.assertRaisesRegex(
         ValueError,
@@ -270,7 +269,7 @@ def test_device_put_concrete_layout(self):
     arr = jax.device_put(np_inp, s)
 
     compiled = jax.jit(
-        lambda x: x * 2).lower(arr, _out_layouts=layout.AUTO).compile()
+        lambda x: x * 2).lower(arr, _out_layouts=DLL.AUTO).compile()
     col = compiled._output_layouts()
 
     out = jax.device_put(np_inp, col)
@@ -283,12 +282,12 @@ def test_device_put_concrete_layout(self):
   def test_device_put_non_concrete_layout_error(self):
     np_inp = np.arange(16).reshape(8, 2)
 
-    l1 = Layout(layout.AUTO, SingleDeviceSharding(jax.devices()[0]))
+    l1 = Layout(DLL.AUTO, SingleDeviceSharding(jax.devices()[0]))
     with self.assertRaisesRegex(
         ValueError, 'sharding and device_local_layout.*should be concrete'):
       jax.device_put(np_inp, l1)
 
-    l2 = Layout(layout.AUTO, None)
+    l2 = Layout(DLL.AUTO, None)
     with self.assertRaisesRegex(
         ValueError, 'sharding and device_local_layout.*should be concrete'):
       jax.device_put(np_inp, l2)