diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 3baf2ad5b9..de20d380d4 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -5,7 +5,8 @@ from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, from tinygrad.device import Buffer, Compiled, Device, MultiBuffer from tinygrad.dtype import DType, dtypes from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite -from tinygrad.engine.realize import capturing, Estimates, compile_linear, link_linear, run_linear, graph_cache, estimate_uop, get_runtime +from tinygrad.renderer import Estimates +from tinygrad.engine.realize import capturing, compile_linear, link_linear, run_linear, graph_cache, estimate_uop, get_runtime from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_outs_ins from tinygrad.schedule.memory import memory_plan_rewrite, _collect_bufs from tinygrad.nn.state import get_parameters diff --git a/tinygrad/nn/onnx.py b/tinygrad/nn/onnx.py index fc55573453..6e537c1bcb 100644 --- a/tinygrad/nn/onnx.py +++ b/tinygrad/nn/onnx.py @@ -2,12 +2,12 @@ from typing import Any, Sequence, cast, Literal, NamedTuple, Generator import dataclasses, functools, io, math, types, warnings, pathlib, sys, os, struct, enum from tinygrad.nn.state import TensorIO -from tinygrad.tensor import Tensor, _broadcast_shape +from tinygrad.tensor import Tensor from tinygrad.mixin import ReductionStr from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN, Context from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate, least_upper_dtype, DTYPES_DICT from tinygrad.device import Device -from tinygrad.uop.ops import sint +from tinygrad.uop.ops import sint, _broadcast_shape # ***** protobuf definitions ****** class WireType(enum.IntEnum): diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 12a6a2dc85..1e0e4d1527 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -2,8 +2,8 @@ from __future__ import annotations import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, contextlib, sys, subprocess, struct assert sys.platform != 'win32' from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler -from tinygrad.dtype import dtypes -from tinygrad.uop.ops import Ops, UOp, AddrSpace +from tinygrad.dtype import dtypes, AddrSpace +from tinygrad.uop.ops import Ops, UOp from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG, suppress_finalizing, Target from tinygrad.renderer.cstyle import ClangRenderer from tinygrad.runtime.autogen import libc, qcom_dsp diff --git a/tinygrad/runtime/support/nv/ip.py b/tinygrad/runtime/support/nv/ip.py index f943ae1c08..75a208d05b 100644 --- a/tinygrad/runtime/support/nv/ip.py +++ b/tinygrad/runtime/support/nv/ip.py @@ -3,7 +3,8 @@ import ctypes, time, array, struct, itertools, dataclasses from typing import cast, Any from tinygrad.runtime.autogen import nv, nv_570 as nv_gpu, pci from tinygrad.helpers import lo32, hi32, DEBUG, round_up, round_down, fetch_fw, wait_cond, ceildiv -from tinygrad.runtime.support.system import System, MMIOInterface +from tinygrad.runtime.support.system import System +from tinygrad.runtime.support.hcq import MMIOInterface from tinygrad.runtime.support.elf import elf_loader @dataclasses.dataclass(frozen=True) diff --git a/tinygrad/runtime/support/nv/nvdev.py b/tinygrad/runtime/support/nv/nvdev.py index cadc599cca..dc5121b0c1 100644 --- a/tinygrad/runtime/support/nv/nvdev.py +++ b/tinygrad/runtime/support/nv/nvdev.py @@ -4,7 +4,8 @@ from tinygrad.helpers import getenv, DEBUG, getbits, round_up from tinygrad.runtime.autogen import pci from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace from tinygrad.runtime.support.nv.ip import NV_FLCN, NV_FLCN_COT, NV_GSP -from tinygrad.runtime.support.system import PCIDevice, MMIOInterface +from tinygrad.runtime.support.system import PCIDevice +from tinygrad.runtime.support.hcq import MMIOInterface NV_DEBUG = getenv("NV_DEBUG", 0) diff --git a/tinygrad/schedule/allreduce.py b/tinygrad/schedule/allreduce.py index 7e210ce512..542520ff0d 100644 --- a/tinygrad/schedule/allreduce.py +++ b/tinygrad/schedule/allreduce.py @@ -1,6 +1,7 @@ import functools, itertools from tinygrad.helpers import all_int, prod, DEBUG, RING, ALL2ALL, getenv -from tinygrad.uop.ops import UOp, Invalid +from tinygrad.uop.ops import UOp +from tinygrad.dtype import Invalid # *** allreduce implementation *** def handle_allreduce(buf:UOp, red:UOp) -> UOp|None: