Fix mypy

Author: saulshanabrook

python/egglog/conversion.py

@@ -180,7 +180,7 @@ def resolve_literal(
     # Try all parent types as well, if we are converting from a Python type
     for arg_type_instance in arg_type.__mro__ if isinstance(arg_type, type) else [arg_type]:
         try:
-            fn = CONVERSIONS[(cast(TypeName | type, arg_type_instance), tp_name)][1]
+            fn = CONVERSIONS[(arg_type_instance, tp_name)][1]
         except KeyError:
             continue
         break

python/egglog/exp/array_api_jit.py

@@ -14,8 +14,6 @@ def jit(fn: X) -> X:
     """
     Jit compiles a function
     """
-    from IPython.display import SVG
-
     # 1. Create variables for each of the two args in the functions
     sig = inspect.signature(fn)
     arg1, arg2 = sig.parameters.keys()
@@ -25,13 +23,12 @@ def jit(fn: X) -> X:
         egraph.register(res)
         egraph.run(array_api_numba_schedule)
         res_optimized = egraph.extract(res)
-        svg = SVG(egraph.graphviz_svg(split_primitive_outputs=True, n_inline_leaves=3))
+        egraph.display(split_primitive_outputs=True, n_inline_leaves=3)
 
     egraph = EGraph()
     fn_program = ndarray_function_two(res_optimized, NDArray.var(arg1), NDArray.var(arg2))
     egraph.register(fn_program)
     egraph.run(array_api_program_gen_schedule)
     fn = cast(X, egraph.eval(fn_program.py_object))
-    fn.egraph = svg  # type: ignore[attr-defined]
     fn.expr = res_optimized  # type: ignore[attr-defined]
     return fn