egglog 6.0.0__cp310-none-win_amd64.whl → 6.1.0__cp310-none-win_amd64.whl

egglog/exp/array_api.py

@@ -25,8 +25,7 @@ if TYPE_CHECKING:
 numbers.Integral.register(RuntimeExpr)
 
 array_api_ruleset = ruleset()
-# Saturate doesn't work for some reason, so just set a high number here
-array_api_schedule = array_api_ruleset * 100000
+array_api_schedule = array_api_ruleset.saturate()
 
 
 class Boolean(Expr):
@@ -35,14 +34,11 @@ class Boolean(Expr):
         return try_evaling(self, self.bool)
 
     @property
-    def bool(self) -> Bool:
-        ...
+    def bool(self) -> Bool: ...
 
-    def __or__(self, other: Boolean) -> Boolean:
-        ...
+    def __or__(self, other: Boolean) -> Boolean: ...
 
-    def __and__(self, other: Boolean) -> Boolean:
-        ...
+    def __and__(self, other: Boolean) -> Boolean: ...
 
 
 TRUE = constant("TRUE", Boolean)
@@ -92,22 +88,18 @@ class IsDtypeKind(Expr):
     NULL: ClassVar[IsDtypeKind]
 
     @classmethod
-    def string(cls, s: StringLike) -> IsDtypeKind:
-        ...
+    def string(cls, s: StringLike) -> IsDtypeKind: ...
 
     @classmethod
-    def dtype(cls, d: DType) -> IsDtypeKind:
-        ...
+    def dtype(cls, d: DType) -> IsDtypeKind: ...
 
     @method(cost=10)
-    def __or__(self, other: IsDtypeKind) -> IsDtypeKind:
-        ...
+    def __or__(self, other: IsDtypeKind) -> IsDtypeKind: ...
 
 
 # TODO: Make kind more generic to support tuples.
 @function
-def isdtype(dtype: DType, kind: IsDtypeKind) -> Boolean:
-    ...
+def isdtype(dtype: DType, kind: IsDtypeKind) -> Boolean: ...
 
 
 converter(DType, IsDtypeKind, lambda x: IsDtypeKind.dtype(x))
@@ -143,17 +135,13 @@ def _isdtype(d: DType, k1: IsDtypeKind, k2: IsDtypeKind):
 
 
 class Int(Expr):
-    def __init__(self, value: i64Like) -> None:
-        ...
+    def __init__(self, value: i64Like) -> None: ...
 
-    def __invert__(self) -> Int:
-        ...
+    def __invert__(self) -> Int: ...
 
-    def __lt__(self, other: Int) -> Boolean:
-        ...
+    def __lt__(self, other: Int) -> Boolean: ...
 
-    def __le__(self, other: Int) -> Boolean:
-        ...
+    def __le__(self, other: Int) -> Boolean: ...
 
     def __eq__(self, other: Int) -> Boolean:  # type: ignore[override]
         ...
@@ -164,93 +152,64 @@ class Int(Expr):
     def __ne__(self, other: Int) -> bool:  # type: ignore[override]
         return not (self == other)
 
-    def __gt__(self, other: Int) -> Boolean:
-        ...
+    def __gt__(self, other: Int) -> Boolean: ...
 
-    def __ge__(self, other: Int) -> Boolean:
-        ...
+    def __ge__(self, other: Int) -> Boolean: ...
 
-    def __add__(self, other: Int) -> Int:
-        ...
+    def __add__(self, other: Int) -> Int: ...
 
-    def __sub__(self, other: Int) -> Int:
-        ...
+    def __sub__(self, other: Int) -> Int: ...
 
-    def __mul__(self, other: Int) -> Int:
-        ...
+    def __mul__(self, other: Int) -> Int: ...
 
-    def __truediv__(self, other: Int) -> Int:
-        ...
+    def __truediv__(self, other: Int) -> Int: ...
 
-    def __floordiv__(self, other: Int) -> Int:
-        ...
+    def __floordiv__(self, other: Int) -> Int: ...
 
-    def __mod__(self, other: Int) -> Int:
-        ...
+    def __mod__(self, other: Int) -> Int: ...
 
-    def __divmod__(self, other: Int) -> Int:
-        ...
+    def __divmod__(self, other: Int) -> Int: ...
 
-    def __pow__(self, other: Int) -> Int:
-        ...
+    def __pow__(self, other: Int) -> Int: ...
 
-    def __lshift__(self, other: Int) -> Int:
-        ...
+    def __lshift__(self, other: Int) -> Int: ...
 
-    def __rshift__(self, other: Int) -> Int:
-        ...
+    def __rshift__(self, other: Int) -> Int: ...
 
-    def __and__(self, other: Int) -> Int:
-        ...
+    def __and__(self, other: Int) -> Int: ...
 
-    def __xor__(self, other: Int) -> Int:
-        ...
+    def __xor__(self, other: Int) -> Int: ...
 
-    def __or__(self, other: Int) -> Int:
-        ...
+    def __or__(self, other: Int) -> Int: ...
 
-    def __radd__(self, other: Int) -> Int:
-        ...
+    def __radd__(self, other: Int) -> Int: ...
 
-    def __rsub__(self, other: Int) -> Int:
-        ...
+    def __rsub__(self, other: Int) -> Int: ...
 
-    def __rmul__(self, other: Int) -> Int:
-        ...
+    def __rmul__(self, other: Int) -> Int: ...
 
-    def __rmatmul__(self, other: Int) -> Int:
-        ...
+    def __rmatmul__(self, other: Int) -> Int: ...
 
-    def __rtruediv__(self, other: Int) -> Int:
-        ...
+    def __rtruediv__(self, other: Int) -> Int: ...
 
-    def __rfloordiv__(self, other: Int) -> Int:
-        ...
+    def __rfloordiv__(self, other: Int) -> Int: ...
 
-    def __rmod__(self, other: Int) -> Int:
-        ...
+    def __rmod__(self, other: Int) -> Int: ...
 
-    def __rpow__(self, other: Int) -> Int:
-        ...
+    def __rpow__(self, other: Int) -> Int: ...
 
-    def __rlshift__(self, other: Int) -> Int:
-        ...
+    def __rlshift__(self, other: Int) -> Int: ...
 
-    def __rrshift__(self, other: Int) -> Int:
-        ...
+    def __rrshift__(self, other: Int) -> Int: ...
 
-    def __rand__(self, other: Int) -> Int:
-        ...
+    def __rand__(self, other: Int) -> Int: ...
 
-    def __rxor__(self, other: Int) -> Int:
-        ...
+    def __rxor__(self, other: Int) -> Int: ...
 
-    def __ror__(self, other: Int) -> Int:
-        ...
+    def __ror__(self, other: Int) -> Int: ...
 
     @property
-    def i64(self) -> i64:
-        ...
+    def i64(self) -> i64: ...
 
     @method(preserve=True)
     def __int__(self) -> int:
@@ -305,31 +264,23 @@ converter(int, Int, lambda x: Int(x))
 
 
 class Float(Expr):
-    def __init__(self, value: f64Like) -> None:
-        ...
+    def __init__(self, value: f64Like) -> None: ...
 
-    def abs(self) -> Float:
-        ...
+    def abs(self) -> Float: ...
 
     @classmethod
-    def rational(cls, r: Rational) -> Float:
-        ...
+    def rational(cls, r: Rational) -> Float: ...
 
     @classmethod
-    def from_int(cls, i: Int) -> Float:
-        ...
+    def from_int(cls, i: Int) -> Float: ...
 
-    def __truediv__(self, other: Float) -> Float:
-        ...
+    def __truediv__(self, other: Float) -> Float: ...
 
-    def __mul__(self, other: Float) -> Float:
-        ...
+    def __mul__(self, other: Float) -> Float: ...
 
-    def __add__(self, other: Float) -> Float:
-        ...
+    def __add__(self, other: Float) -> Float: ...
 
-    def __sub__(self, other: Float) -> Float:
-        ...
+    def __sub__(self, other: Float) -> Float: ...
 
 
 converter(float, Float, lambda x: Float(x))
@@ -357,14 +308,11 @@ def _float(f: f64, f2: f64, i: i64, r: Rational, r1: Rational):
 class TupleInt(Expr):
     EMPTY: ClassVar[TupleInt]
 
-    def __init__(self, head: Int) -> None:
-        ...
+    def __init__(self, head: Int) -> None: ...
 
-    def __add__(self, other: TupleInt) -> TupleInt:
-        ...
+    def __add__(self, other: TupleInt) -> TupleInt: ...
 
-    def length(self) -> Int:
-        ...
+    def length(self) -> Int: ...
 
     @method(preserve=True)
     def __len__(self) -> int:
@@ -374,11 +322,9 @@ class TupleInt(Expr):
     def __iter__(self) -> Iterator[Int]:
         return iter(self[Int(i)] for i in range(len(self)))
 
-    def __getitem__(self, i: Int) -> Int:
-        ...
+    def __getitem__(self, i: Int) -> Int: ...
 
-    def product(self) -> Int:
-        ...
+    def product(self) -> Int: ...
 
 
 converter(
@@ -413,8 +359,7 @@ class OptionalInt(Expr):
     none: ClassVar[OptionalInt]
 
     @classmethod
-    def some(cls, value: Int) -> OptionalInt:
-        ...
+    def some(cls, value: Int) -> OptionalInt: ...
 
 
 converter(type(None), OptionalInt, lambda _: OptionalInt.none)
@@ -427,8 +372,7 @@ class Slice(Expr):
         start: OptionalInt = OptionalInt.none,
         stop: OptionalInt = OptionalInt.none,
         step: OptionalInt = OptionalInt.none,
-    ) -> None:
-        ...
+    ) -> None: ...
 
 
 converter(
@@ -443,12 +387,10 @@ class MultiAxisIndexKeyItem(Expr):
     NONE: ClassVar[MultiAxisIndexKeyItem]
 
     @classmethod
-    def int(cls, i: Int) -> MultiAxisIndexKeyItem:
-        ...
+    def int(cls, i: Int) -> MultiAxisIndexKeyItem: ...
 
     @classmethod
-    def slice(cls, slice: Slice) -> MultiAxisIndexKeyItem:
-        ...
+    def slice(cls, slice: Slice) -> MultiAxisIndexKeyItem: ...
 
 
 converter(type(...), MultiAxisIndexKeyItem, lambda _: MultiAxisIndexKeyItem.ELLIPSIS)
@@ -458,13 +400,11 @@ converter(Slice, MultiAxisIndexKeyItem, MultiAxisIndexKeyItem.slice)
 
 
 class MultiAxisIndexKey(Expr):
-    def __init__(self, item: MultiAxisIndexKeyItem) -> None:
-        ...
+    def __init__(self, item: MultiAxisIndexKeyItem) -> None: ...
 
     EMPTY: ClassVar[MultiAxisIndexKey]
 
-    def __add__(self, other: MultiAxisIndexKey) -> MultiAxisIndexKey:
-        ...
+    def __add__(self, other: MultiAxisIndexKey) -> MultiAxisIndexKey: ...
 
 
 converter(
@@ -492,12 +432,10 @@ class IndexKey(Expr):
     ELLIPSIS: ClassVar[IndexKey]
 
     @classmethod
-    def int(cls, i: Int) -> IndexKey:
-        ...
+    def int(cls, i: Int) -> IndexKey: ...
 
     @classmethod
-    def slice(cls, slice: Slice) -> IndexKey:
-        ...
+    def slice(cls, slice: Slice) -> IndexKey: ...
 
     # Disabled until we support late binding
     # @classmethod
@@ -505,8 +443,7 @@ class IndexKey(Expr):
     #     ...
 
     @classmethod
-    def multi_axis(cls, key: MultiAxisIndexKey) -> IndexKey:
-        ...
+    def multi_axis(cls, key: MultiAxisIndexKey) -> IndexKey: ...
 
 
 converter(type(...), IndexKey, lambda _: IndexKey.ELLIPSIS)
@@ -515,8 +452,7 @@ converter(Slice, IndexKey, IndexKey.slice)
 converter(MultiAxisIndexKey, IndexKey, IndexKey.multi_axis)
 
 
-class Device(Expr):
-    ...
+class Device(Expr): ...
 
 
 ALL_INDICES: TupleInt = constant("ALL_INDICES", TupleInt)
@@ -528,28 +464,21 @@ ALL_INDICES: TupleInt = constant("ALL_INDICES", TupleInt)
 
 class Value(Expr):
     @classmethod
-    def int(cls, i: Int) -> Value:
-        ...
+    def int(cls, i: Int) -> Value: ...
 
     @classmethod
-    def float(cls, f: Float) -> Value:
-        ...
+    def float(cls, f: Float) -> Value: ...
 
     @classmethod
-    def bool(cls, b: Boolean) -> Value:
-        ...
+    def bool(cls, b: Boolean) -> Value: ...
 
-    def isfinite(self) -> Boolean:
-        ...
+    def isfinite(self) -> Boolean: ...
 
-    def __lt__(self, other: Value) -> Value:
-        ...
+    def __lt__(self, other: Value) -> Value: ...
 
-    def __truediv__(self, other: Value) -> Value:
-        ...
+    def __truediv__(self, other: Value) -> Value: ...
 
-    def astype(self, dtype: DType) -> Value:
-        ...
+    def astype(self, dtype: DType) -> Value: ...
 
     # TODO: Add all operations
 
@@ -560,12 +489,10 @@ class Value(Expr):
         """
 
     @property
-    def to_bool(self) -> Boolean:
-        ...
+    def to_bool(self) -> Boolean: ...
 
     @property
-    def to_int(self) -> Int:
-        ...
+    def to_int(self) -> Int: ...
 
     @property
     def to_truthy_value(self) -> Value:
@@ -600,20 +527,15 @@ def _value(i: Int, f: Float, b: Boolean):
 class TupleValue(Expr):
     EMPTY: ClassVar[TupleValue]
 
-    def __init__(self, head: Value) -> None:
-        ...
+    def __init__(self, head: Value) -> None: ...
 
-    def __add__(self, other: TupleValue) -> TupleValue:
-        ...
+    def __add__(self, other: TupleValue) -> TupleValue: ...
 
-    def length(self) -> Int:
-        ...
+    def length(self) -> Int: ...
 
-    def __getitem__(self, i: Int) -> Value:
-        ...
+    def __getitem__(self, i: Int) -> Value: ...
 
-    def includes(self, value: Value) -> Boolean:
-        ...
+    def includes(self, value: Value) -> Boolean: ...
 
 
 converter(
@@ -661,41 +583,34 @@ def possible_values(values: Value) -> TupleValue:
 
 
 class NDArray(Expr):
-    def __init__(self, py_array: PyObject) -> None:
-        ...
+    def __init__(self, py_array: PyObject) -> None: ...
 
     @method(cost=200)
     @classmethod
-    def var(cls, name: StringLike) -> NDArray:
-        ...
+    def var(cls, name: StringLike) -> NDArray: ...
 
     @method(preserve=True)
     def __array_namespace__(self, api_version: object = None) -> ModuleType:
         return sys.modules[__name__]
 
     @property
-    def ndim(self) -> Int:
-        ...
+    def ndim(self) -> Int: ...
 
     @property
-    def dtype(self) -> DType:
-        ...
+    def dtype(self) -> DType: ...
 
     @property
-    def device(self) -> Device:
-        ...
+    def device(self) -> Device: ...
 
     @property
-    def shape(self) -> TupleInt:
-        ...
+    def shape(self) -> TupleInt: ...
 
     @method(preserve=True)
     def __bool__(self) -> bool:
         return bool(self.to_value().to_bool)
 
     @property
-    def size(self) -> Int:
-        ...
+    def size(self) -> Int: ...
 
     @method(preserve=True)
     def __len__(self) -> int:
@@ -706,17 +621,13 @@ class NDArray(Expr):
         for i in range(len(self)):
             yield self[IndexKey.int(Int(i))]
 
-    def __getitem__(self, key: IndexKey) -> NDArray:
-        ...
+    def __getitem__(self, key: IndexKey) -> NDArray: ...
 
-    def __setitem__(self, key: IndexKey, value: NDArray) -> None:
-        ...
+    def __setitem__(self, key: IndexKey, value: NDArray) -> None: ...
 
-    def __lt__(self, other: NDArray) -> NDArray:
-        ...
+    def __lt__(self, other: NDArray) -> NDArray: ...
 
-    def __le__(self, other: NDArray) -> NDArray:
-        ...
+    def __le__(self, other: NDArray) -> NDArray: ...
 
     def __eq__(self, other: NDArray) -> NDArray:  # type: ignore[override]
         ...
@@ -725,109 +636,77 @@ class NDArray(Expr):
     # def __ne__(self, other: NDArray) -> NDArray:  # type: ignore[override]
     #     ...
 
-    def __gt__(self, other: NDArray) -> NDArray:
-        ...
+    def __gt__(self, other: NDArray) -> NDArray: ...
 
-    def __ge__(self, other: NDArray) -> NDArray:
-        ...
+    def __ge__(self, other: NDArray) -> NDArray: ...
 
-    def __add__(self, other: NDArray) -> NDArray:
-        ...
+    def __add__(self, other: NDArray) -> NDArray: ...
 
-    def __sub__(self, other: NDArray) -> NDArray:
-        ...
+    def __sub__(self, other: NDArray) -> NDArray: ...
 
-    def __mul__(self, other: NDArray) -> NDArray:
-        ...
+    def __mul__(self, other: NDArray) -> NDArray: ...
 
-    def __matmul__(self, other: NDArray) -> NDArray:
-        ...
+    def __matmul__(self, other: NDArray) -> NDArray: ...
 
-    def __truediv__(self, other: NDArray) -> NDArray:
-        ...
+    def __truediv__(self, other: NDArray) -> NDArray: ...
 
-    def __floordiv__(self, other: NDArray) -> NDArray:
-        ...
+    def __floordiv__(self, other: NDArray) -> NDArray: ...
 
-    def __mod__(self, other: NDArray) -> NDArray:
-        ...
+    def __mod__(self, other: NDArray) -> NDArray: ...
 
-    def __divmod__(self, other: NDArray) -> NDArray:
-        ...
+    def __divmod__(self, other: NDArray) -> NDArray: ...
 
-    def __pow__(self, other: NDArray) -> NDArray:
-        ...
+    def __pow__(self, other: NDArray) -> NDArray: ...
 
-    def __lshift__(self, other: NDArray) -> NDArray:
-        ...
+    def __lshift__(self, other: NDArray) -> NDArray: ...
 
-    def __rshift__(self, other: NDArray) -> NDArray:
-        ...
+    def __rshift__(self, other: NDArray) -> NDArray: ...
 
-    def __and__(self, other: NDArray) -> NDArray:
-        ...
+    def __and__(self, other: NDArray) -> NDArray: ...
 
-    def __xor__(self, other: NDArray) -> NDArray:
-        ...
+    def __xor__(self, other: NDArray) -> NDArray: ...
 
-    def __or__(self, other: NDArray) -> NDArray:
-        ...
+    def __or__(self, other: NDArray) -> NDArray: ...
 
-    def __radd__(self, other: NDArray) -> NDArray:
-        ...
+    def __radd__(self, other: NDArray) -> NDArray: ...
 
-    def __rsub__(self, other: NDArray) -> NDArray:
-        ...
+    def __rsub__(self, other: NDArray) -> NDArray: ...
 
-    def __rmul__(self, other: NDArray) -> NDArray:
-        ...
+    def __rmul__(self, other: NDArray) -> NDArray: ...
 
-    def __rmatmul__(self, other: NDArray) -> NDArray:
-        ...
+    def __rmatmul__(self, other: NDArray) -> NDArray: ...
 
-    def __rtruediv__(self, other: NDArray) -> NDArray:
-        ...
+    def __rtruediv__(self, other: NDArray) -> NDArray: ...
 
-    def __rfloordiv__(self, other: NDArray) -> NDArray:
-        ...
+    def __rfloordiv__(self, other: NDArray) -> NDArray: ...
 
-    def __rmod__(self, other: NDArray) -> NDArray:
-        ...
+    def __rmod__(self, other: NDArray) -> NDArray: ...
 
-    def __rpow__(self, other: NDArray) -> NDArray:
-        ...
+    def __rpow__(self, other: NDArray) -> NDArray: ...
 
-    def __rlshift__(self, other: NDArray) -> NDArray:
-        ...
+    def __rlshift__(self, other: NDArray) -> NDArray: ...
 
-    def __rrshift__(self, other: NDArray) -> NDArray:
-        ...
+    def __rrshift__(self, other: NDArray) -> NDArray: ...
 
-    def __rand__(self, other: NDArray) -> NDArray:
-        ...
+    def __rand__(self, other: NDArray) -> NDArray: ...
 
-    def __rxor__(self, other: NDArray) -> NDArray:
-        ...
+    def __rxor__(self, other: NDArray) -> NDArray: ...
 
-    def __ror__(self, other: NDArray) -> NDArray:
-        ...
+    def __ror__(self, other: NDArray) -> NDArray: ...
 
     @classmethod
-    def scalar(cls, value: Value) -> NDArray:
-        ...
+    def scalar(cls, value: Value) -> NDArray: ...
 
-    def to_value(self) -> Value:
-        ...
+    def to_value(self) -> Value: ...
 
     @property
-    def T(self) -> NDArray:  # noqa: N802
+    def T(self) -> NDArray:
         """
         https://data-apis.org/array-api/2022.12/API_specification/generated/array_api.array.T.html#array_api.array.T
         """
 
     @classmethod
-    def vector(cls, values: TupleValue) -> NDArray:
-        ...
+    def vector(cls, values: TupleValue) -> NDArray: ...
 
     def index(self, indices: TupleInt) -> Value:
         """
@@ -878,14 +757,11 @@ def _ndarray(x: NDArray, b: Boolean, f: Float, fi1: f64, fi2: f64):
 class TupleNDArray(Expr):
     EMPTY: ClassVar[TupleNDArray]
 
-    def __init__(self, head: NDArray) -> None:
-        ...
+    def __init__(self, head: NDArray) -> None: ...
 
-    def __add__(self, other: TupleNDArray) -> TupleNDArray:
-        ...
+    def __add__(self, other: TupleNDArray) -> TupleNDArray: ...
 
-    def length(self) -> Int:
-        ...
+    def length(self) -> Int: ...
 
     @method(preserve=True)
     def __len__(self) -> int:
@@ -895,8 +771,7 @@ class TupleNDArray(Expr):
     def __iter__(self) -> Iterator[NDArray]:
         return iter(self[Int(i)] for i in range(len(self)))
 
-    def __getitem__(self, i: Int) -> NDArray:
-        ...
+    def __getitem__(self, i: Int) -> NDArray: ...
 
 
 converter(
@@ -928,8 +803,7 @@ class OptionalBool(Expr):
     none: ClassVar[OptionalBool]
 
     @classmethod
-    def some(cls, value: Boolean) -> OptionalBool:
-        ...
+    def some(cls, value: Boolean) -> OptionalBool: ...
 
 
 converter(type(None), OptionalBool, lambda _: OptionalBool.none)
@@ -940,8 +814,7 @@ class OptionalDType(Expr):
     none: ClassVar[OptionalDType]
 
     @classmethod
-    def some(cls, value: DType) -> OptionalDType:
-        ...
+    def some(cls, value: DType) -> OptionalDType: ...
 
 
 converter(type(None), OptionalDType, lambda _: OptionalDType.none)
@@ -952,8 +825,7 @@ class OptionalDevice(Expr):
     none: ClassVar[OptionalDevice]
 
     @classmethod
-    def some(cls, value: Device) -> OptionalDevice:
-        ...
+    def some(cls, value: Device) -> OptionalDevice: ...
 
 
 converter(type(None), OptionalDevice, lambda _: OptionalDevice.none)
@@ -964,8 +836,7 @@ class OptionalTupleInt(Expr):
     none: ClassVar[OptionalTupleInt]
 
     @classmethod
-    def some(cls, value: TupleInt) -> OptionalTupleInt:
-        ...
+    def some(cls, value: TupleInt) -> OptionalTupleInt: ...
 
 
 converter(type(None), OptionalTupleInt, lambda _: OptionalTupleInt.none)
@@ -976,12 +847,10 @@ class IntOrTuple(Expr):
     none: ClassVar[IntOrTuple]
 
     @classmethod
-    def int(cls, value: Int) -> IntOrTuple:
-        ...
+    def int(cls, value: Int) -> IntOrTuple: ...
 
     @classmethod
-    def tuple(cls, value: TupleInt) -> IntOrTuple:
-        ...
+    def tuple(cls, value: TupleInt) -> IntOrTuple: ...
 
 
 converter(Int, IntOrTuple, IntOrTuple.int)
@@ -992,8 +861,7 @@ class OptionalIntOrTuple(Expr):
     none: ClassVar[OptionalIntOrTuple]
 
     @classmethod
-    def some(cls, value: IntOrTuple) -> OptionalIntOrTuple:
-        ...
+    def some(cls, value: IntOrTuple) -> OptionalIntOrTuple: ...
 
 
 converter(type(None), OptionalIntOrTuple, lambda _: OptionalIntOrTuple.none)
@@ -1001,8 +869,9 @@ converter(IntOrTuple, OptionalIntOrTuple, OptionalIntOrTuple.some)
 
 
 @function
-def asarray(a: NDArray, dtype: OptionalDType = OptionalDType.none, copy: OptionalBool = OptionalBool.none) -> NDArray:
-    ...
+def asarray(
+    a: NDArray, dtype: OptionalDType = OptionalDType.none, copy: OptionalBool = OptionalBool.none
+) -> NDArray: ...
 
 
 @array_api_ruleset.register
@@ -1012,8 +881,7 @@ def _assarray(a: NDArray, d: OptionalDType, ob: OptionalBool):
 
 
 @function
-def isfinite(x: NDArray) -> NDArray:
-    ...
+def isfinite(x: NDArray) -> NDArray: ...
 
 
 @function
@@ -1032,8 +900,7 @@ def _sum(x: NDArray, y: NDArray, v: Value, dtype: DType):
 
 
 @function
-def reshape(x: NDArray, shape: TupleInt, copy: OptionalBool = OptionalBool.none) -> NDArray:
-    ...
+def reshape(x: NDArray, shape: TupleInt, copy: OptionalBool = OptionalBool.none) -> NDArray: ...
 
 
 # @function
@@ -1075,8 +942,7 @@ def reshape(x: NDArray, shape: TupleInt, copy: OptionalBool = OptionalBool.none)
 
 
 @function
-def unique_values(x: NDArray) -> NDArray:
-    ...
+def unique_values(x: NDArray) -> NDArray: ...
 
 
 @array_api_ruleset.register
@@ -1087,8 +953,7 @@ def _unique_values(x: NDArray):
 
 
 @function
-def concat(arrays: TupleNDArray, axis: OptionalInt = OptionalInt.none) -> NDArray:
-    ...
+def concat(arrays: TupleNDArray, axis: OptionalInt = OptionalInt.none) -> NDArray: ...
 
 
 @array_api_ruleset.register
@@ -1099,8 +964,7 @@ def _concat(x: NDArray):
 
 
 @function
-def astype(x: NDArray, dtype: DType) -> NDArray:
-    ...
+def astype(x: NDArray, dtype: DType) -> NDArray: ...
 
 
 @array_api_ruleset.register
@@ -1113,9 +977,8 @@ def _astype(x: NDArray, dtype: DType, i: i64):
     ]
 
 
-@function(cost=500)
-def unique_counts(x: NDArray) -> TupleNDArray:
-    ...
+@function
+def unique_counts(x: NDArray) -> TupleNDArray: ...
 
 
 @array_api_ruleset.register
@@ -1131,23 +994,19 @@ def _unique_counts(x: NDArray, c: NDArray, tv: TupleValue, v: Value, dtype: DTyp
 
 
 @function
-def square(x: NDArray) -> NDArray:
-    ...
+def square(x: NDArray) -> NDArray: ...
 
 
 @function
-def any(x: NDArray) -> NDArray:
-    ...
+def any(x: NDArray) -> NDArray: ...
 
 
 @function(egg_fn="ndarray-abs")
-def abs(x: NDArray) -> NDArray:
-    ...
+def abs(x: NDArray) -> NDArray: ...
 
 
 @function(egg_fn="ndarray-log")
-def log(x: NDArray) -> NDArray:
-    ...
+def log(x: NDArray) -> NDArray: ...
 
 
 @array_api_ruleset.register
@@ -1157,9 +1016,8 @@ def _abs(f: Float):
     ]
 
 
-@function(cost=100)
-def unique_inverse(x: NDArray) -> TupleNDArray:
-    ...
+@function
+def unique_inverse(x: NDArray) -> TupleNDArray: ...
 
 
 @array_api_ruleset.register
@@ -1174,29 +1032,24 @@ def _unique_inverse(x: NDArray, i: Int):
 @function
 def zeros(
     shape: TupleInt, dtype: OptionalDType = OptionalDType.none, device: OptionalDevice = OptionalDevice.none
-) -> NDArray:
-    ...
+) -> NDArray: ...
 
 
 @function
-def expand_dims(x: NDArray, axis: Int = Int(0)) -> NDArray:
-    ...
+def expand_dims(x: NDArray, axis: Int = Int(0)) -> NDArray: ...
 
 
-@function(cost=100000)
-def mean(x: NDArray, axis: OptionalIntOrTuple = OptionalIntOrTuple.none, keepdims: Boolean = FALSE) -> NDArray:
-    ...
+@function
+def mean(x: NDArray, axis: OptionalIntOrTuple = OptionalIntOrTuple.none, keepdims: Boolean = FALSE) -> NDArray: ...
 
 
 # TODO: Possibly change names to include modules.
 @function(egg_fn="ndarray-sqrt")
-def sqrt(x: NDArray) -> NDArray:
-    ...
+def sqrt(x: NDArray) -> NDArray: ...
 
 
-@function(cost=100000)
-def std(x: NDArray, axis: OptionalIntOrTuple = OptionalIntOrTuple.none) -> NDArray:
-    ...
+@function
+def std(x: NDArray, axis: OptionalIntOrTuple = OptionalIntOrTuple.none) -> NDArray: ...
 
 
 linalg = sys.modules[__name__]
@@ -1485,13 +1338,11 @@ def _size(x: NDArray):
 
 
 @overload
-def try_evaling(expr: Expr, prim_expr: i64) -> int:
-    ...
+def try_evaling(expr: Expr, prim_expr: i64) -> int: ...
 
 
 @overload
-def try_evaling(expr: Expr, prim_expr: Bool) -> bool:
-    ...
+def try_evaling(expr: Expr, prim_expr: Bool) -> bool: ...
 
 
 def try_evaling(expr: Expr, prim_expr: i64 | Bool) -> int | bool: