Better error reporting

einx/op/index.py

@@ -354,6 +354,24 @@ def after_stage2(exprs1, exprs2):
     )[: len(op[0]) + 1]
     exprs_in, expr_out = exprs[: len(op[0])], exprs[len(op[0])]
 
+    if update:
+        # Check that all axes in first input expression also appear in output expression
+        axes_in = {
+            axis.name
+            for axis in exprs_in[0].all()
+            if isinstance(axis, einx.expr.stage3.Axis)
+        }
+        axes_out = {
+            axis.name
+            for axis in expr_out.all()
+            if isinstance(axis, einx.expr.stage3.Axis)
+        }
+        if not axes_in.issubset(axes_out):
+            raise ValueError(
+                f"Output expression does not contain all axes from first input expression: "
+                f"{axes_in - axes_out}"
+            )
+
     return exprs_in, expr_out