Change jax.core.DropVar to be a non-singleton.

Previously jax.core.DropVar was a singleton value (jax.core.dropvar) whose type was always jax.core.AbstractUnit. However, this type is misleading: a DropVar is an equation output, and typically we would expect it to have an array type. In particular, the unit type confuses new-style translation rules that expect to use the output aval on an equation as part of the lowering logic.

Instead, change DropVar to be a non-singleton subclass of Var instead with a flexible choice of aval.

PiperOrigin-RevId: 404071001

spinoffs/oryx/oryx/core/interpreters/propagate.py

@@ -164,7 +164,7 @@ def write(self, var: VarOrLiteral, cell: Cell) -> Cell:
     if isinstance(var, jax_core.Literal):
       return cell
     cur_cell = self.read(var)
-    if var is jax_core.dropvar:
+    if isinstance(var, jax_core.DropVar):
       return cur_cell
     self.env[var] = cur_cell.join(cell)
     return self.env[var]