fix typing FlopCounter.flops can be sint (#3646)

This commit is contained in:
chenyu
2024-03-07 12:49:17 -05:00
committed by GitHub
parent 906cc3a69b
commit 0cef284aac

View File

@@ -4,7 +4,7 @@ import functools, hashlib
from enum import Enum, auto
from tinygrad.helpers import prod, dedup
from tinygrad.dtype import dtypes, DType
from tinygrad.shape.symbolic import Variable
from tinygrad.shape.symbolic import Variable, sint
from dataclasses import dataclass
# these are the llops your accelerator must implement, along with toCpu
@@ -75,7 +75,7 @@ class LazyOp:
class FlopCounter:
shape: Tuple[int, ...]
dtype: DType
flops: int
flops: sint
mem: Dict[int, int]
@property
def mem_estimate(self): return sum(self.mem.values())