Misc bpr typing fixes, update baseline

Author: inducer

.basedpyright/baseline.json

@@ -12,3 +12,12 @@ Generic target support
 ----------------------
 
 .. automodule:: pytato.target
+
+References
+----------
+
+.. currentmodule:: cl_array
+
+.. class:: Allocator
+
+    See :class:`pyopencl.array.Allocator`.

doc/codegen.rst

@@ -178,6 +178,30 @@ ignore = [
 ]
 
 
+[[tool.basedpyright.executionEnvironments]]
+root = "test"
+reportUnknownArgumentType = "none"
+reportUnknownVariableType = "none"
+reportMissingParameterType = "none"
+reportAttributeAccessIssue = "none"
+reportMissingImports = "none"
+reportArgumentType = "hint"
+reportUnknownMemberType = "hint"
+reportUnknownParameterType = "hint"
+reportAny = "none"
+
+[[tool.basedpyright.executionEnvironments]]
+root = "examples"
+reportUnknownArgumentType = "none"
+reportUnknownVariableType = "none"
+reportMissingParameterType = "none"
+reportAttributeAccessIssue = "none"
+reportMissingImports = "none"
+reportArgumentType = "hint"
+reportUnknownMemberType = "hint"
+reportUnknownParameterType = "hint"
+reportAny = "none"
+
 [tool.typos.default]
 extend-ignore-re = [
   "(?Rm)^.*(#|//)\\s*spellchecker:\\s*disable-line$"

pyproject.toml

@@ -60,7 +60,7 @@
 
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable, Iterator, Mapping, Sequence
+    from collections.abc import Iterator, Mapping, Sequence
 
     from pymbolic.typing import ArithmeticExpression, Expression, Integer
 
@@ -338,8 +338,8 @@ def _get_val_in_bset(bset: isl.BasicSet, idim: int) -> ScalarExpression:
     return aff_to_expr(aff)
 
 
-def solve_constraints(variables: Iterable[str],
-                      parameters: Iterable[str],
+def solve_constraints(variables: Sequence[str],
+                      parameters: Sequence[str],
                       constraints: Sequence[tuple[ArithmeticExpression,
                                                   ArithmeticExpression]],
 
@@ -520,10 +520,10 @@ def extend_bindings_with_shape_inference(knl: lp.LoopKernel,
 
     # }}}
 
-    solutions = solve_constraints(variables={_lp_var_to_global_namespace(var)
-                                             for var in lp_size_params},
-                                  parameters={_pt_var_to_global_namespace(var.name)
-                                              for var in pt_size_params},
+    solutions = solve_constraints(variables=list({_lp_var_to_global_namespace(var)
+                                             for var in lp_size_params}),
+                                  parameters=list({_pt_var_to_global_namespace(var.name)
+                                              for var in pt_size_params}),
                                   constraints=constraints)
 
     as_pt_size_param = EvaluationMapper({_pt_var_to_global_namespace(arg.name): arg

pytato/loopy.py

@@ -130,7 +130,7 @@ def _normalize_pad_width(
 
 
 def pad(array: Array,
-        pad_width: Integer | Sequence[Integer],
+        pad_width: int | Sequence[int],
         mode: str = "constant",
         **kwargs: Any) -> Array:
     r"""

pytato/pad.py

@@ -109,7 +109,8 @@ def map_foreign(self, expr: Any, depth: int) -> str:
         elif isinstance(expr, frozenset | set):
             return "{" + ", ".join(self.rec(el, depth) for el in expr) + "}"
         elif isinstance(expr, np.dtype):
-            return f"'{expr.name}'"
+            dtype_name = cast("str", expr.name)
+            return f"'{dtype_name}'"
         else:
             return repr(expr)
 

pytato/stringifier.py

@@ -67,6 +67,7 @@
     from collections.abc import Callable, Mapping
 
     import pyopencl
+    import pyopencl.array as cl_array
 
 
 class ImplSubstitution(ImplementationStrategy):
@@ -176,7 +177,7 @@ def with_transformed_program(self, f: Callable[[loopy.TranslationUnit],
     def _get_processed_bound_arguments(
                 self,
                 queue: pyopencl.CommandQueue,
-                allocator: Callable[[int], pyopencl.MemoryObject] | None,
+                allocator: cl_array.Allocator | None,
             ) -> Mapping[str, Any]:
         import pyopencl.array as cla
 
@@ -216,7 +217,8 @@ def all_bound_args_on_host(self) -> bool:
                    for arg in self.bound_arguments.values())
 
     def __call__(self, queue: pyopencl.CommandQueue,  # type: ignore[no-untyped-def,no-any-unimported]
-                 allocator=None, wait_for=None, out_host: bool | None = None,
+                 allocator: cl_array.Allocator | None = None,
+                 wait_for: pyopencl.WaitList = None, out_host: bool | None = None,
                  **kwargs: Any) -> Any:
         """Convenience function for launching a :mod:`pyopencl` computation."""
 

pytato/target/loopy/__init__.py

@@ -78,6 +78,7 @@
 
 
 if TYPE_CHECKING:
+    from ast import _ConstantValue  # pyright: ignore[reportPrivateUsage]
     from collections.abc import Callable, Iterable
 
     from pytato.target.python import BoundPythonProgram, NumpyLikePythonTarget
@@ -127,6 +128,10 @@ def first_true(iterable: Iterable[T], default: T,
     return next(filter(pred, iterable), default)
 
 
+def _constant(value: object) -> ast.Constant:
+    return ast.Constant(cast("_ConstantValue", value))
+
+
 def _is_slice_trivial(slice_: NormalizedSlice,
                       dim: ShapeComponent) -> bool:
     """
@@ -217,26 +222,26 @@ def _rec_ary_or_constant(e: ArrayOrScalar) -> ast.expr:
                     # generates code like: `np.float64("nan")`.
                     return ast.Call(
                         func=ast.Attribute(value=ast.Name(self.numpy),
-                                           attr=e_np.dtype.name),
-                        args=[ast.Constant(value="nan")],
+                                           attr=cast("str", e_np.dtype.name)),
+                        args=[_constant(value="nan")],
                         keywords=[])
                 else:
-                    return ast.Constant(e)
+                    return _constant(e)
 
         if isinstance(hlo, FullOp):
             if hlo.fill_value == 1:
                 if expr.dtype == np.dtype(float):
                     rhs = ast.Call(
                         ast.Attribute(ast.Name(self.numpy_backend),
                                       "ones"),
-                        args=[ast.Tuple(elts=[ast.Constant(d)
+                        args=[ast.Tuple(elts=[_constant(d)
                                               for d in expr.shape])],
                         keywords=[])
                 else:
                     rhs = ast.Call(
                         ast.Attribute(ast.Name(self.numpy_backend),
                                       "ones"),
-                        args=[ast.Tuple(elts=[ast.Constant(d)
+                        args=[ast.Tuple(elts=[_constant(d)
                                               for d in expr.shape])],
                         keywords=[ast.keyword(
                             arg="dtype",
@@ -248,14 +253,14 @@ def _rec_ary_or_constant(e: ArrayOrScalar) -> ast.expr:
                     rhs = ast.Call(
                         ast.Attribute(ast.Name(self.numpy_backend),
                                       "zeros"),
-                        args=[ast.Tuple(elts=[ast.Constant(d)
+                        args=[ast.Tuple(elts=[_constant(d)
                                               for d in expr.shape])],
                         keywords=[])
                 else:
                     rhs = ast.Call(
                         ast.Attribute(ast.Name(self.numpy_backend),
                                       "zeros"),
-                        args=[ast.Tuple(elts=[ast.Constant(d)
+                        args=[ast.Tuple(elts=[_constant(d)
                                               for d in expr.shape])],
                         keywords=[ast.keyword(
                             arg="dtype",
@@ -266,7 +271,7 @@ def _rec_ary_or_constant(e: ArrayOrScalar) -> ast.expr:
                 rhs = ast.Call(
                     ast.Attribute(ast.Name(self.numpy_backend),
                                   "full"),
-                    args=[ast.Tuple(elts=[ast.Constant(d)
+                    args=[ast.Tuple(elts=[_constant(d)
                                           for d in expr.shape]),
                           _rec_ary_or_constant(hlo.fill_value),
                           ],
@@ -324,7 +329,7 @@ def _rec_ary_or_constant(e: ArrayOrScalar) -> ast.expr:
             rhs = ast.Call(ast.Attribute(ast.Name(self.numpy_backend),
                                          "broadcast_to"),
                            args=[ast.Name(self.rec(hlo.x)),
-                                 ast.Tuple(elts=[ast.Constant(d)
+                                 ast.Tuple(elts=[_constant(d)
                                                  for d in expr.shape])],
                            keywords=[])
         elif isinstance(hlo, ReduceOp):
@@ -339,9 +344,9 @@ def _rec_ary_or_constant(e: ArrayOrScalar) -> ast.expr:
             else:
                 if len(hlo.axes) == 1:
                     axis, = hlo.axes.keys()
-                    axis_ast: ast.expr = ast.Constant(axis)
+                    axis_ast: ast.expr = _constant(axis)
                 else:
-                    axis_ast = ast.Tuple(elts=[ast.Constant(e)
+                    axis_ast = ast.Tuple(elts=[_constant(e)
                                                for e in sorted(hlo.axes.keys())])
                 rhs = ast.Call(ast.Attribute(ast.Name(self.numpy_backend),
                                              np_fn_name),
@@ -366,7 +371,7 @@ def map_stack(self, expr: Stack) -> str:
                        args=[ast.List([ast.Name(id_)
                                        for id_ in rec_ids])],
                        keywords=[ast.keyword(arg="axis",
-                                             value=ast.Constant(expr.axis))])
+                                             value=_constant(expr.axis))])
 
         return self._record_line_and_return_lhs(lhs, rhs)
 
@@ -379,7 +384,7 @@ def map_concatenate(self, expr: Concatenate) -> str:
                        args=[ast.List([ast.Name(id_)
                                        for id_ in rec_ids])],
                        keywords=[ast.keyword(arg="axis",
-                                             value=ast.Constant(expr.axis))])
+                                             value=_constant(expr.axis))])
 
         return self._record_line_and_return_lhs(lhs, rhs)
 
@@ -389,9 +394,9 @@ def map_roll(self, expr: Roll) -> str:
                        args=[ast.Name(self.rec(expr.array)),
                              ],
                        keywords=[ast.keyword(arg="shift",
-                                             value=ast.Constant(expr.shift)),
+                                             value=_constant(expr.shift)),
                                  ast.keyword(arg="axis",
-                                             value=ast.Constant(expr.axis))])
+                                             value=_constant(expr.axis))])
 
         return self._record_line_and_return_lhs(lhs, rhs)
 
@@ -404,7 +409,7 @@ def map_axis_permutation(self, expr: AxisPermutation) -> str:
                            args=[ast.Name(self.rec(expr.array))],
                            keywords=[ast.keyword(
                                arg="axes",
-                               value=ast.List(elts=[ast.Constant(a)
+                               value=ast.List(elts=[_constant(a)
                                                     for a in expr.axis_permutation]))
                                      ])
 
@@ -427,7 +432,7 @@ def _map_index_base(self, expr: IndexBase) -> str:
 
         def _rec_idx(idx: IndexExpr, dim: ShapeComponent) -> ast.expr:
             if isinstance(idx, int):
-                return ast.Constant(idx)
+                return _constant(idx)
             elif isinstance(idx, NormalizedSlice):
                 step = idx.step if idx.step != 1 else None
                 if idx.step > 0:
@@ -458,13 +463,13 @@ class SliceKwargs(TypedDict):
                 kwargs: SliceKwargs = {}
                 if step is not None:
                     assert isinstance(step, int)
-                    kwargs["step"] = ast.Constant(step)
+                    kwargs["step"] = _constant(step)
                 if start is not None:
                     assert isinstance(start, int)
-                    kwargs["lower"] = ast.Constant(start)
+                    kwargs["lower"] = _constant(start)
                 if stop is not None:
                     assert isinstance(stop, int)
-                    kwargs["upper"] = ast.Constant(stop)
+                    kwargs["upper"] = _constant(stop)
 
                 return ast.Slice(**kwargs)
             else:
@@ -500,7 +505,7 @@ def map_einsum(self, expr: Einsum) -> str:
         lhs = self.vng("_pt_tmp")
         args = [ast.Name(self.rec(arg)) for arg in expr.args]
         rhs = ast.Call(ast.Attribute(ast.Name(self.numpy_backend), "einsum"),
-                        args=[ast.Constant(get_einsum_specification(expr)),
+                        args=[_constant(get_einsum_specification(expr)),
                               *args],
                        keywords=[],
                        )
@@ -513,7 +518,7 @@ def map_reshape(self, expr: Reshape) -> str:
             raise NotImplementedError("Non-integral reshapes.")
         rhs = ast.Call(ast.Attribute(ast.Name(self.numpy_backend), "reshape"),
                         args=[ast.Name(self.rec(expr.array)),
-                              ast.Tuple(elts=[ast.Constant(d)
+                              ast.Tuple(elts=[_constant(d)
                                               for d in expr.shape])],
                        keywords=[],
                        )
@@ -530,7 +535,7 @@ def map_dict_of_named_arrays(self, expr: DictOfNamedArrays) -> str:
         keys: list[expr_t | None] = []
         values: list[expr_t] = []
         for name, subexpr in sorted(expr._data.items()):
-            keys.append(ast.Constant(name))
+            keys.append(_constant(name))
             values.append(ast.Name(self.rec(subexpr)))
 
         rhs = ast.Dict(keys=keys, values=values)

pytato/target/python/numpy_like.py

@@ -2366,7 +2366,7 @@ def _get_data_dedup_cache_key(self, ary: DataInterface) -> CacheKeyT:
 
             if isinstance(ary, CLArray):
                 base_data = ary.base_data
-                if isinstance(ary.base_data, MemoryObjectHolder):
+                if isinstance(base_data, MemoryObjectHolder):
                     ptr = base_data.int_ptr
                 elif SVMPointer is not None and isinstance(base_data, SVMPointer):
                     ptr = base_data.svm_ptr

pytato/transform/__init__.py

@@ -423,8 +423,7 @@ def _is_non_negative(expr: ShapeComponent) -> Bool:
                                       for expr in InputGatherer()(expr)})
     aff = ShapeToISLExpressionMapper(space)(expr)
 
-    # type-ignore because islpy is not typed yet
-    return (aff.ge_set(aff * 0) >= _get_size_params_assumptions_bset(space))  # type: ignore[no-any-return]
+    return aff.ge_set(aff * 0) >= _get_size_params_assumptions_bset(space).to_set()
 
 
 def _is_non_positive(expr: ShapeComponent) -> Bool:
@@ -485,7 +484,9 @@ def _normalize_slice(slice_: slice,
 
 
 def _normalized_slice_len(slice_: NormalizedSlice) -> ShapeComponent:
-    start, stop, step = slice_.start, slice_.stop, slice_.step
+    start, stop, step = cast(
+                "tuple[int, int, int]",
+                (slice_.start, slice_.stop, slice_.step))
 
     if step > 0:
         if _is_non_negative(stop - start):