Repository: google/jax
Pull request #21717 (2024-06-06 21:38:46)

rename `Specialized` to `Traced` (and `specialize` to `trace`)
rename `Specialized` to `Traced` (and `specialize` to `trace`)



Patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
index aa4579475..26f5d15bb 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1813,10 +1813,10 @@ def _cpp_pmap(
 
   @api_boundary
   def lower(*args, **kwargs):
-    return specialize(*args, **kwargs).lower()
+    return trace(*args, **kwargs).lower()
 
   @api_boundary
-  def specialize(*args, **kwargs):
+  def trace(*args, **kwargs):
     lowering_parameters = kwargs.pop(
         '_experimental_lowering_parameters', mlir.LoweringParameters())
     p = _prepare_pmap(
@@ -1847,11 +1847,11 @@ def _cpp_pmap(
         shards=shards,
         pci=pci)
     args_info = stages.make_args_info(p.in_tree, abstract_args, donate_tuple)
-    return stages.Specialized(closed_jaxpr, args_info, p.flat_fun.__name__,
-                              p.out_tree(), lower_callable)
+    return stages.Traced(closed_jaxpr, args_info, p.flat_fun.__name__,
+                         p.out_tree(), lower_callable)
 
   pmap_f.lower = lower
-  pmap_f.specialize = specialize
+  pmap_f.trace = trace
 
   return pmap_f
 
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index a139227ad..50fb4bfc1 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -477,16 +477,16 @@ def _make_jit_wrapper(jit_info: PjitInfo):
 
   @api_boundary
   def lower(*args, **kwargs):
-    specialized = specialize(*args, **kwargs)
+    traced = trace(*args, **kwargs)
     try:
-      return specialized.lower()
+      return traced.lower()
     except pxla.DeviceAssignmentMismatchError as e:
       fails, = e.args
       fun = jit_info.fun
       fun_name = getattr(fun, '__qualname__',
                          getattr(fun, '__name__', str(fun)))
       msg = _device_assignment_mismatch_error(
-          fun_name, fails, specialized._args_flat, 'jit', specialized._arg_names)
+          fun_name, fails, traced._args_flat, 'jit', traced._arg_names)
       raise ValueError(msg) from None
 
   @api_boundary
@@ -499,7 +499,7 @@ def _make_jit_wrapper(jit_info: PjitInfo):
     return tree_unflatten(out_tree, out)
 
   @api_boundary
-  def specialize(*args, **kwargs) -> stages.Specialized:
+  def trace(*args, **kwargs) -> stages.Traced:
     lowering_parameters = kwargs.pop(
         '_experimental_lowering_parameters', mlir.LoweringParameters())
 
@@ -511,13 +511,13 @@ def _make_jit_wrapper(jit_info: PjitInfo):
     args_info = stages.make_args_info(in_tree, jaxpr.in_avals, donate_argnums)
     lower_callable = partial(_resolve_and_lower, args_flat, **params,
                              lowering_parameters=lowering_parameters)
-    return stages.Specialized(jaxpr, args_info, params["name"], out_tree,
-                              lower_callable, args_flat, arg_names)
+    return stages.Traced(jaxpr, args_info, params["name"], out_tree,
+                         lower_callable, args_flat, arg_names)
 
   wrapped = _cpp_pjit(jit_info)
   wrapped.lower = lower
   wrapped.eval_shape = eval_shape
-  wrapped.specialize = specialize
+  wrapped.trace = trace
   return wrapped
 
 
@@ -711,7 +711,7 @@ class JitWrapped(stages.Wrapped):
     """See ``jax.eval_shape``."""
     raise NotImplementedError
 
-  def specialize(self, *args, **kwargs) -> stages.Specialized:
+  def trace(self, *args, **kwargs) -> stages.Traced:
     raise NotImplementedError
 
 
diff --git a/jax/_src/stages.py b/jax/_src/stages.py
index 90ad765a9..03ce788ab 100644
--- a/jax/_src/stages.py
+++ b/jax/_src/stages.py
@@ -425,7 +425,7 @@ class CompiledCallParams(NamedTuple):
   out_tree: tree_util.PyTreeDef
 
 
-class Specialized(Stage):
+class Traced(Stage):
   __slots__ = ["jaxpr", "args_info", "fun_name", "_out_tree", "_lower_callable",
                "_args_flat", "_arg_names"]
 
diff --git a/jax/experimental/export/_export.py b/jax/experimental/export/_export.py
index 8b13e8d63..c0e325382 100644
--- a/jax/experimental/export/_export.py
+++ b/jax/experimental/export/_export.py
@@ -452,7 +452,7 @@ def export(fun_jax: Callable,
       # an error if the lowered function contains non-replicated sharding annotations.
       wrapped_fun_jax = jax.jit(fun_jax)
 
-    has_specialize = hasattr(wrapped_fun_jax, "specialize")
+    has_trace = hasattr(wrapped_fun_jax, "trace")
 
     if lowering_platforms is not None:
       actual_lowering_platforms = tuple(lowering_platforms)
@@ -475,15 +475,15 @@ def export(fun_jax: Callable,
               self_descr=f"current (from {shape_poly.args_kwargs_path_to_str(symbolic_scope[1])}) ",
               other_descr=shape_poly.args_kwargs_path_to_str(k_path))
 
-    if has_specialize:
-      specialized = wrapped_fun_jax.specialize(
+    if has_trace:
+      traced = wrapped_fun_jax.trace(
           *args_specs, **kwargs_specs,
           _experimental_lowering_parameters=mlir.LoweringParameters(
             platforms=actual_lowering_platforms,
             for_export=True,
           ))
-      jaxpr, fun_name = specialized.jaxpr, specialized.fun_name
-      lowered = specialized.lower()
+      jaxpr, fun_name = traced.jaxpr, traced.fun_name
+      lowered = traced.lower()
     else:
       lowered = wrapped_fun_jax.lower(
           *args_specs, **kwargs_specs,
diff --git a/jax/stages.py b/jax/stages.py
index 280ddd5c4..6ffc3144c 100644
--- a/jax/stages.py
+++ b/jax/stages.py
@@ -31,5 +31,5 @@ from jax._src.stages import (
   Wrapped as Wrapped,
   ArgInfo as ArgInfo,
   OutInfo as OutInfo,
-  Specialized as Specialized,
+  Traced as Traced,
 )


Test patch:
diff --git a/tests/export_test.py b/tests/export_test.py
index 1383ca44e..73dee5a20 100644
--- a/tests/export_test.py
+++ b/tests/export_test.py
@@ -559,8 +559,8 @@ class JaxExportTest(jtu.JaxTestCase):
             x,
             _experimental_lowering_parameters=_experimental_lowering_parameters)
 
-      def specialize(self, x, _experimental_lowering_parameters=None):
-        return jax.jit(self.__call__).specialize(
+      def trace(self, x, _experimental_lowering_parameters=None):
+        return jax.jit(self.__call__).trace(
             x,
             _experimental_lowering_parameters=_experimental_lowering_parameters)
 
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 92afbe455..2a75a2719 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4159,25 +4159,25 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertEqual(out_info[0].sharding, None)
     self.assertEqual(out_info[1].sharding, None)
 
-  def test_jit_specialize(self):
+  def test_jit_trace(self):
     def f(x):
       return x * 2
 
-    specialized = jax.jit(f).specialize(jnp.arange(8, dtype=jnp.int32))
-    self.assertLen(specialized.jaxpr.eqns, 1)
-    self.assertEqual(jax.tree.structure(specialized.out_info).num_leaves, 1)
-    self.assertEqual(specialized.out_info.shape, (8,))
-    self.assertEqual(specialized.out_info.dtype, jnp.int32)
+    traced = jax.jit(f).trace(jnp.arange(8, dtype=jnp.int32))
+    self.assertLen(traced.jaxpr.eqns, 1)
+    self.assertEqual(jax.tree.structure(traced.out_info).num_leaves, 1)
+    self.assertEqual(traced.out_info.shape, (8,))
+    self.assertEqual(traced.out_info.dtype, jnp.int32)
     # one for args, one for kwargs (though kwargs is empty)
-    self.assertLen(specialized.in_avals, 2)
-    self.assertLen(specialized.in_avals[0], 1)
-    self.assertLen(specialized.in_avals[1], 0)  # empty kwarg
+    self.assertLen(traced.in_avals, 2)
+    self.assertLen(traced.in_avals[0], 1)
+    self.assertLen(traced.in_avals[1], 0)  # empty kwarg
 
-  def test_jit_specialize_lower_and_compile(self):
+  def test_jit_trace_lower_and_compile(self):
     def f(x):
       return x * 2
 
-    lowered = jax.jit(f).specialize(jnp.arange(8)).lower()
+    lowered = jax.jit(f).trace(jnp.arange(8)).lower()
     self.assertEqual(lowered.args_info[0][0].shape, (8,))
 
     compiled = lowered.compile()