Add to_cotangent_aval to HiType and use it by default in bwd pass

PiperOrigin-RevId: 845470915

Author: yashk2810

jax/_src/api.py

@@ -2332,9 +2332,7 @@ def _vjp_check_ct_avals(cts, primal_avals):
   # TODO(mattjj): improve this error  by flattening with keys in the first place
   for ct, aval in zip(cts, primal_avals):
     ct_aval = typeof(ct)
-    ct_aval_expected = (
-        aval.to_cotangent_aval() if hasattr(aval, 'to_cotangent_aval') else
-        aval.to_tangent_aval())
+    ct_aval_expected = aval.to_cotangent_aval()
     if (not core.typecompat(ct_aval, ct_aval_expected) and
         not _temporary_dtype_exception(ct_aval, ct_aval_expected)):
       raise ValueError(

jax/_src/core.py

@@ -1654,6 +1654,9 @@ class AbstractValue:
   def to_tangent_aval(self):
     raise NotImplementedError("must override")
 
+  def to_cotangent_aval(self):
+    raise NotImplementedError("must override")
+
   # TODO(dougalm): deprecate this alias
   def at_least_vspace(self):
     return self.to_tangent_aval()
@@ -2619,6 +2622,7 @@ def accum_grad_in_ref(x):
 class AbstractToken(AbstractValue):
   def str_short(self, short_dtypes=False, mesh_axis_types=False): return 'Tok'
   def to_tangent_aval(self): return self
+  def to_cotangent_aval(self): return self
 abstract_token: AbstractToken = AbstractToken()
 
 # Singleton shaped array used by all abstract tokens when shape/dtype is needed.

jax/_src/hijax.py

@@ -92,9 +92,16 @@ def raise_val(self, *lo_vals: LoVal) -> HiVal:
   # autodiff interface
   def to_tangent_aval(self) -> HiType:
     assert False, "must override"
+
+  # Subclasses should override if the cotangent type is a function of primal
+  # type. For example, CT unreduced = reduced and vice-versa.
+  def to_cotangent_aval(self) -> HiType:
+    return self.to_tangent_aval()
+
   # the next two are required if this type is itself a tangent type
   def vspace_zero(self) -> HiVal:
     assert False, "must override"
+
   def vspace_add(self, x: HiVal, y: HiVal) -> HiVal:
     assert False, "must override"
 
@@ -127,6 +134,11 @@ def update_from_loval(self, state: QDD, val: HiVal, *lo_vals: LoVal) -> None:
   def to_tangent_aval(self) -> HiType:
     assert False, "must override"
 
+  # Subclasses should override if the cotangent type is a function of primal
+  # type. For example, CT unreduced = reduced and vice-versa.
+  def to_cotangent_aval(self) -> HiType:
+    return self.to_tangent_aval()
+
 def register_hitype(val_cls, typeof_fn) -> None:
   core.pytype_aval_mappings[val_cls] = typeof_fn
   dtypes.canonicalize_value_handlers[val_cls] = lambda x: x