Remove type parameters. Type checker doesn't like them

Author: dougalm

jax/_src/interpreters/partial_eval.py

@@ -2292,9 +2292,9 @@ def _jvp_jaxpr_zeros(f, store, in_zeros, zero_avals, *primal_tangent_avals):
 @weakref_lru_cache
 def trace_to_jaxpr(
     fun: Callable,
-    in_avals: FlatTree[AbstractValue | core.AvalQDD],  # (args, kwargs) pair
+    in_avals: FlatTree,  # (args, kwargs) pair
     debug_info: core.DebugInfo
-) -> tuple[ClosedJaxpr, PyTreeDef]:
+) -> tuple[ClosedJaxpr, FlatTree]:
   config.enable_checks.value and debug_info.assert_arg_names(len(in_avals))
   parent_trace = core.trace_ctx.trace
   trace = DynamicJaxprTrace(debug_info, parent_trace=parent_trace)

jax/_src/lax/control_flow/loops.py

@@ -84,7 +84,7 @@ def _stack(arrs: Sequence[Array], axis: int=0) -> Array:
 
 def _promote_weak_typed_input(
     in_val:Any, in_aval:AbstractValue, out_aval:AbstractValue
-    ) -> tuple[AbstractValue, bool]:
+    ) -> tuple[Any, bool]:
   if getattr(in_aval, 'weak_type', False) and not core.typematch(in_aval, out_aval):
     new_dtype = dtypes.result_type(in_val, out_aval)
     return lax.convert_element_type(in_val, new_dtype), True
@@ -228,7 +228,7 @@ def scan(f, init, xs, length=None):
     return carry, stacked_y
 
   if config.mutable_array_checks.value:
-    check_no_aliased_ref_args(lambda: dbg_body, list(args), list(args_avals))
+    check_no_aliased_ref_args(lambda: dbg_body, list(args_avals), list(args))
 
   x_avals = xs_avals.map(lambda aval: core.mapped_aval(length, 0, aval))
   def _create_jaxpr(carry_avals):
@@ -252,6 +252,8 @@ def _create_jaxpr(carry_avals):
   if config.mutable_array_checks.value:
     _check_no_aliased_closed_over_refs(dbg_body, consts, list(args))
   carry_out_avals, ys_avals = out_avals.unpack()
+  if len(carry_out_avals) != len(init_avals):
+    _check_carry_type('scan body', f, init_avals, carry_out_avals)
   init, changed = init.map3(
      _promote_weak_typed_input,
      init_avals, carry_out_avals).unzip2()
@@ -277,7 +279,7 @@ def _create_jaxpr(carry_avals):
   if unroll < 0:
     raise ValueError("`unroll` must be a `bool` or a non-negative `int`.")
 
-  args_flat = (*init.vals, *xs.vals)
+  args_flat = [*init.vals, *xs.vals]
 
   # If the body forwards an input carry to an output carry, that input is
   # read-only and can be moved to be a const. Doing so can lead to efficiency
@@ -381,18 +383,18 @@ def _check_carry_type(name, body_fun, in_carry, out_carry):
                            if p else 'the input carry')
   if in_carry.tree != out_carry.tree:
     try:
-      out_carry = out_carry.unflatten()
+      out_carry_unflat = out_carry.unflatten()
     except:
-      out_carry = None
+      out_carry_unflat = None
 
-    if out_carry is None:
+    if out_carry_unflat is None:
       differences = (f'the input tree structure is:\n{in_carry.tree}\n' +
                      f'the output tree structure is:\n{out_carry.tree}\n')
     else:
       diffs = [f'{component(path)} is a {thing1} but the corresponding component '
                f'of the carry output is a {thing2}, so {explanation}'
                for path, thing1, thing2, explanation
-               in equality_errors(in_carry, out_carry)]
+               in equality_errors(in_carry.unflatten(), out_carry.unflatten())]
       if len(diffs) == 0:
         return  # the trees may have different aux data, but structures are same
       elif len(diffs) == 1:
@@ -1709,7 +1711,7 @@ def _create_jaxpr(init_avals):
 
   cond_dbg = api_util.debug_info("while_cond", cond_fun, (init_val,), {})
   body_dbg = api_util.debug_info("while_body", body_fun, (init_val,), {})
-  init_val = FlatTree.flatten(init_val)
+  init_val = FlatTree.flatten(init_val)  # type: ignore
   init_aval = init_val.map(core.get_aval)
 
   # The body input and output avals must match exactly. However, we want to account for
@@ -1718,6 +1720,10 @@ def _create_jaxpr(init_avals):
   # To do this, we compute the jaxpr in two passes: first with the raw inputs, and if
   # necessary, a second time with modified init values.
   cond_jaxpr, body_jaxpr, body_out_avals = _create_jaxpr(init_aval)
+  if len(body_out_avals) != len(init_aval):
+    _check_carry_type('while_loop body', body_fun, init_aval, body_out_avals)
+    assert False, "shouldn't get here"
+
   init_val, changed = init_val.map3(
       _promote_weak_typed_input,
       init_aval, body_out_avals).unzip2()
@@ -1749,7 +1755,7 @@ def _create_jaxpr(init_avals):
   _, keep_cond_carry = split_list(keep_cond, [len(cond_consts)])
   move_to_const = _map(operator.not_, keep_cond_carry)
 
-  init_vals = list(init_val)
+  init_vals = list(init_val)  # type: ignore
   if any(move_to_const):
     cond_jaxpr = pe.close_jaxpr(cond_jaxpr_)
     body_jaxpr = pe.prune_closed_jaxpr_outputs(

jax/_src/tree_util.py

@@ -34,10 +34,6 @@
 traceback_util.register_exclusion(__file__)
 
 T = TypeVar("T")
-T1 = TypeVar("T1")
-T2 = TypeVar("T2")
-T3 = TypeVar("T3")
-T4 = TypeVar("T4")
 Typ = TypeVar("Typ", bound=type[Any])
 H = TypeVar("H", bound=Hashable)
 
@@ -1357,20 +1353,18 @@ class FlatTree:
        the tuple-returning function would change the tree structure and `unzip`
        wouldn't be able to recover it.
   """
-  def __init__(self, vals:Sequence[T], treedef:PyTreeDef):
+  def __init__(self, vals:Sequence, treedef:PyTreeDef):
     assert isinstance(treedef, pytree.PyTreeDef)
     self.tree = treedef
-    self.vals = list(vals)
+    self.vals = tuple(vals)
 
-  def map(self, f:Callable[[T1], T2]) -> FlatTree[T2]:
+  def map(self, f:Callable) -> FlatTree:
     ans_vals = []
     for x in self.vals:
       ans_vals.append(f(x))
     return FlatTree(ans_vals, self.tree)
 
-  def map2(
-      self:FlatTree[T1], f:Callable[[T1, T2], T3],
-      t2:FlatTree[T2]) -> FlatTree[T3]:
+  def map2(self:FlatTree, f:Callable, t2:FlatTree) -> FlatTree:
 
     n = len(self)
     assert len(t2) == n
@@ -1380,19 +1374,18 @@ def map2(
     return FlatTree(ans_vals, self.tree)
 
   def map3(
-      self:FlatTree[T1], f:Callable[[T1, T2, T3], T4],
-      t2:FlatTree[T2], t3:FlatTree[T3]) -> FlatTree[T4]:
+      self:FlatTree, f:Callable, t2:FlatTree, t3:FlatTree) -> FlatTree:
     n = len(self)
     assert len(t2) == n and len(t3) == n
     ans_vals = []
     for x1, x2, x3 in zip(self.vals, t2.vals, t3.vals):
       ans_vals.append(f(x1, x2, x3))
     return FlatTree(ans_vals, self.tree)
 
-  def zip(self, t2:FlatTree[T2]) -> FlatTree[tuple[T1, T2]]:
+  def zip(self, t2:FlatTree) -> FlatTree:
     assert False
 
-  def unzip2(self:FlatTree[tuple[T1, T2]]) -> tuple[FlatTree[T1], FlatTree[T2]]:
+  def unzip2(self:FlatTree) -> tuple[FlatTree, FlatTree]:
     ys = []
     zs = []
     for y, z in self.vals:
@@ -1425,7 +1418,7 @@ def pack(tree):
     else:
       assert False
 
-  def unpack(self:FlatTree[tuple]) -> tuple[FlatTree]:
+  def unpack(self:FlatTree) -> tuple[FlatTree, ...]:
     # TODO: this is O(N) not O(1) (with N as the number of leaves). If it
     # becomes a problem we can fix it with a fancier data tree.
     trees = treedef_children(self.tree)
@@ -1444,11 +1437,19 @@ def flatten(tree: PyTree) -> FlatTree:
   def unflatten(self) -> PyTree:
     return tree_unflatten(self.tree, self.vals)
 
-  def update_from_list(self, new_vals:list[T1]) -> FlatTree[T1]:
+  def update_from_list(self, new_vals:list) -> FlatTree:
     return FlatTree(new_vals, self.tree)
 
   def __len__(self):
     return self.tree.num_leaves
 
   def __iter__(self):
     return self.vals.__iter__()
+
+  def __eq__(self, other):
+    return (isinstance(other, FlatTree)
+            and self.vals == other.vals
+            and self.tree == other.tree)
+
+  def __hash__(self):
+    return hash((self.vals, self.tree))