Repository: google/jax
Pull request #21678 (2024-06-05 19:54:33)

Add `specialize` on jax.jit and make it a `Stage`.
Add `specialize` on jax.jit and make it a `Stage`.

Eventually, we should use this in jax.make_jaxpr and delete all the duplicated code.


Fe1777: Sure, here's a structured and detailed GitHub issue description for adding specialization to `jax.jit`.

---

## **Title:** Add Specialize on `jax.jit` and Plan for Subsequent Enhancements

**Type:** Enhancement

**Reported by:** copybara-service

---

### **Description**

This issue aims to enhance the `jax.jit` functionality by introducing specialization. The goal is to make this functionality a dedicated stage and eventually incorporate a `lower` method into it. This enhancement will streamline the code and reduce duplication, particularly in `jax.make_jaxpr`.

### **Objective**

- Introduce specialization in `jax.jit`.
- Plan to create a dedicated stage for specialization.
- Incorporate a `lower` method in the follow-up stages.
- Use this specialization in `jax.make_jaxpr` to eliminate duplicated code.

### **Benefits**

- **Code Optimization**: Reduces code redundancy, especially in `jax.make_jaxpr`.
- **Stepwise Enhancement**: Allows for incremental improvements and easier testing.
- **Improved Functionality**: Enhances the capabilities and flexibility of `jax.jit`.

### **Implementation Plan**

1. **Introduce Specialization**: Add the initial specialization to `jax.jit`.
2. **Create Stage**: In a follow-up change, encapsulate the specialization as its own stage.
3. **Add Lower Method**: Incorporate a `lower` method to the specialization stage.
4. **Refactor jax.make_jaxpr**: Use the newly created specialization in `jax.make_jaxpr` and remove the duplicated code.

### **Example**

**Current Approach:**
```python
import jax

@jax.jit
def f(x):
    return x + x

print(f(3))
```

**Future Approach with Specialization:**
```python
import jax

@jax.jit.spec
def f(x):
    return x + x

f_lowered = jax.jit.lower(f)
print(f_lowered(3))
```

### **Tasks**
- [ ] Implement initial specialization in `jax.jit`.
- [ ] Transition specialization to its own stage.
- [ ] Develop the `lower` method for the specialization stage.
- [ ] Refactor `jax.make_jaxpr` to use the new specialization.
- [ ] Remove duplicated code in `jax.make_jaxpr`.
- [ ] Verify changes by running the test suite.

### **Related Issues/PRs**
- **Related PR:** Add specialize on `jax.jit` ([#21678](https://github.com/google/jax/pull/21678)).

### **Additional Notes**

- Changes should be backward compatible where possible.
- Ensure thorough testing for each step to avoid regressions.

### **Help Wanted**

Developers who are familiar with JAX's codebase and have experience with `jax.jit` and `jax.make_jaxpr` are encouraged to participate. Your contributions will be crucial in making this enhancement robust and efficient.

### **Thank You!**

Thank you for contributing to the enhancement of JAX. Your efforts to improve the functionality and efficiency of this library are highly valued.

---

Feel free to further fine-tune or expand upon this template based on any additional context or specific details you might want to include.

Patch:
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index ae88e75d5..74361dc93 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -509,9 +509,19 @@ def _make_jit_wrapper(jit_info: PjitInfo):
            for x, s in zip(params['jaxpr'].out_avals, out_s)]
     return tree_unflatten(out_tree, out)
 
+  @api_boundary
+  def specialize(*args, **kwargs) -> stages.Specialized:
+    _, _, params, in_tree, out_tree, donated_invars, _, _ = _infer_params(
+        jit_info, args, kwargs)
+    donate_argnums = tuple(i for i, d in enumerate(donated_invars) if d)
+    jaxpr = params['jaxpr']
+    args_info = stages.make_args_info(in_tree, jaxpr.in_avals, donate_argnums)
+    return stages.Specialized(jaxpr, args_info, out_tree)
+
   wrapped = _cpp_pjit(jit_info)
   wrapped.lower = lower
   wrapped.eval_shape = eval_shape
+  wrapped.specialize = specialize
   return wrapped
 
 
@@ -706,6 +716,9 @@ class JitWrapped(stages.Wrapped):
     """See ``jax.eval_shape``."""
     raise NotImplementedError
 
+  def specialize(self, *args, **kwargs) -> stages.Specialized:
+    raise NotImplementedError
+
 
 # in_shardings and out_shardings can't be None as the default value
 # because `None` means that the input is fully replicated.
diff --git a/jax/_src/stages.py b/jax/_src/stages.py
index f23fd1b05..4ac24c261 100644
--- a/jax/_src/stages.py
+++ b/jax/_src/stages.py
@@ -387,7 +387,7 @@ class ArgInfo:
 class OutInfo:
   shape: tuple[int, ...]
   dtype: jax.typing.DTypeLike
-  sharding: jax.sharding.Sharding
+  sharding: jax.sharding.Sharding | None = None
 
 
 class Stage:
@@ -425,6 +425,20 @@ class CompiledCallParams(NamedTuple):
   out_tree: tree_util.PyTreeDef
 
 
+class Specialized(Stage):
+  __slots__ = ["jaxpr", "args_info", "_out_tree"]
+
+  def __init__(self, jaxpr: core.ClosedJaxpr, args_info, out_tree):
+    self.jaxpr = jaxpr
+    self.args_info = args_info
+    self._out_tree = out_tree
+
+  @property
+  def out_info(self):
+    return self._out_tree.unflatten(
+        [OutInfo(o.shape, o.dtype) for o in self.jaxpr.out_avals])
+
+
 class Compiled(Stage):
   """Compiled representation of a function specialized to types/values.
 
diff --git a/jax/stages.py b/jax/stages.py
index e567d9264..280ddd5c4 100644
--- a/jax/stages.py
+++ b/jax/stages.py
@@ -31,4 +31,5 @@ from jax._src.stages import (
   Wrapped as Wrapped,
   ArgInfo as ArgInfo,
   OutInfo as OutInfo,
+  Specialized as Specialized,
 )


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 9c88ee623..58aa8457b 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4177,6 +4177,20 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertEqual(out_info[0].sharding, None)
     self.assertEqual(out_info[1].sharding, None)
 
+  def test_jit_specialize(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)
+    # 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
+
 
 def spec_regex(s):
   return str(s).replace(r"(", r"\(").replace(r")", r"\)")