diff --git a/docs/abstractions2.py b/docs/abstractions2.py index fe9e140a44..34fa4ded39 100644 --- a/docs/abstractions2.py +++ b/docs/abstractions2.py @@ -76,7 +76,7 @@ assert out.as_buffer().cast('I')[0] == 5 print("******** third, the LazyBuffer ***********") -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.engine.realize import run_schedule from tinygrad.engine.schedule import create_schedule diff --git a/docs/developer/developer.md b/docs/developer/developer.md index b4858e371d..4df3fe3724 100644 --- a/docs/developer/developer.md +++ b/docs/developer/developer.md @@ -11,7 +11,7 @@ There is a good [bunch of tutorials](https://mesozoic-egg.github.io/tinygrad-not Everything in [Tensor](../tensor/index.md) is syntactic sugar around [function.py](function.md), where the forwards and backwards passes are implemented for the different functions. There's about 25 of them, implemented using about 20 basic ops. Those basic ops go on to construct a graph of: -::: tinygrad.lazy.LazyBuffer +::: tinygrad.engine.lazy.LazyBuffer options: show_source: false diff --git a/extra/introspection.py b/extra/introspection.py index f3e301aacd..f84f6f6fa7 100644 --- a/extra/introspection.py +++ b/extra/introspection.py @@ -1,7 +1,7 @@ # TODO: move the GRAPH and DEBUG stuff to here import gc from tinygrad.helpers import prod -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.device import Buffer from tinygrad import Tensor, GlobalCounters diff --git a/test/external/external_benchmark_ast.py b/test/external/external_benchmark_ast.py index b7d6dc89a4..9270f8faa1 100644 --- a/test/external/external_benchmark_ast.py +++ b/test/external/external_benchmark_ast.py @@ -6,7 +6,7 @@ from tinygrad import Tensor from tinygrad.codegen.kernel import Kernel from tinygrad.helpers import Context, getenv, to_function_name from tinygrad.engine.schedule import _get_output_groups, _lower_lazybuffer -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.ops import UOp, UOps if __name__ == "__main__": diff --git a/test/external/external_benchmark_multitensor_allreduce.py b/test/external/external_benchmark_multitensor_allreduce.py index 3ad1cc2cc7..b20da756a3 100644 --- a/test/external/external_benchmark_multitensor_allreduce.py +++ b/test/external/external_benchmark_multitensor_allreduce.py @@ -1,6 +1,6 @@ import time from tinygrad import Tensor, Device, GlobalCounters, TinyJit -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.ops import ReduceOps from tinygrad.multi import MultiLazyBuffer, all_reduce from tinygrad.engine.schedule import create_schedule diff --git a/test/external/fuzz_schedule.py b/test/external/fuzz_schedule.py index 752ff994f6..776f2672dd 100644 --- a/test/external/fuzz_schedule.py +++ b/test/external/fuzz_schedule.py @@ -4,7 +4,7 @@ from typing import DefaultDict, Dict, List, Set, Tuple, TypeVar, Union from tinygrad.device import Buffer from tinygrad.engine.realize import CustomOp, capturing, lower_schedule_item from tinygrad.helpers import DEBUG, MULTIOUTPUT, colored, getenv -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.engine.schedule import LBScheduleItem, _graph_schedule, ScheduleItem from tinygrad.ops import MetaOps from tinygrad.tensor import Tensor, _to_np_dtype diff --git a/test/external/process_replay/diff_schedule.py b/test/external/process_replay/diff_schedule.py index 791df044ab..4cf34549c1 100644 --- a/test/external/process_replay/diff_schedule.py +++ b/test/external/process_replay/diff_schedule.py @@ -5,7 +5,7 @@ from typing import DefaultDict, List, Set, Tuple from test.external.process_replay.helpers import print_diff from tinygrad.engine.schedule import LBScheduleItem, ScheduleItem from tinygrad.helpers import CI, DEBUG, Context, ContextVar, colored, diskcache_put, fetch, getenv -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.engine.realize import CompiledRunner, lower_schedule_item CAPTURING_PROCESS_REPLAY = ContextVar("CAPTURING_PROCESS_REPLAY", getenv("RUN_PROCESS_REPLAY")) diff --git a/test/external/process_replay/test_diff_schedule.py b/test/external/process_replay/test_diff_schedule.py index 69c4233d12..9f94faa808 100644 --- a/test/external/process_replay/test_diff_schedule.py +++ b/test/external/process_replay/test_diff_schedule.py @@ -4,7 +4,7 @@ from test.external.process_replay.diff_schedule import CAPTURING_PROCESS_REPLAY, from tinygrad import Tensor, nn from tinygrad.helpers import Context from tinygrad.engine.schedule import _graph_schedule -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer class TestDiffSchedule(unittest.TestCase): def setUp(self): diff --git a/test/test_custom_function.py b/test/test_custom_function.py index aeba369de3..e64d544d72 100644 --- a/test/test_custom_function.py +++ b/test/test_custom_function.py @@ -9,7 +9,7 @@ from tinygrad.dtype import dtypes # *** first, we implement the atan2 op at the lowest level *** # `atan2_gpu` for GPUBuffers and `atan2_cpu` for CPUBuffers -from tinygrad.lazy import Buffer, create_lazybuffer +from tinygrad.engine.lazy import Buffer, create_lazybuffer from tinygrad.device import Device from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.engine.realize import CompiledRunner @@ -32,7 +32,7 @@ def atan2_cpu(ret:Buffer, a:Buffer, b:Buffer): ret.copyin(np.require(np.arctan2( # In general, it is also optional to write a backward function, just your backward pass won't work without it from tinygrad.ops import MetaOps -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.tensor import Function class ATan2(Function): diff --git a/test/test_lazybuffer.py b/test/test_lazybuffer.py index 35b8876a9d..1ca672c92f 100644 --- a/test/test_lazybuffer.py +++ b/test/test_lazybuffer.py @@ -3,7 +3,7 @@ import numpy as np import unittest from tinygrad import Tensor, Device, dtypes from tinygrad.ops import UOps -from tinygrad.lazy import LazyBuffer, MetaOps +from tinygrad.engine.lazy import LazyBuffer, MetaOps from tinygrad.engine.schedule import create_schedule class TestLazyBuffer(unittest.TestCase): diff --git a/test/test_renderer_failures.py b/test/test_renderer_failures.py index 818d2ba4d9..2d49fa5362 100644 --- a/test/test_renderer_failures.py +++ b/test/test_renderer_failures.py @@ -10,7 +10,7 @@ from tinygrad.renderer.cstyle import CStyleLanguage from tinygrad.ops import BinaryOps, UOp, UOps from tinygrad.renderer import Program from tinygrad.tensor import Tensor, _to_np_dtype -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer def _test_uop_result(inputs:List[Tensor], stores:List[UOp], local_size=None): for x in inputs: x.realize() diff --git a/test/test_schedule.py b/test/test_schedule.py index 24047042c5..46d5efc177 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -20,7 +20,7 @@ from tinygrad.codegen.kernel import Kernel, verify_ast from tinygrad.engine.schedule import BUF_LIMIT, create_schedule, reduceop_fusor, st_fixup from tinygrad.engine.realize import CompiledRunner, run_schedule from test.helpers import ast_const, is_dtype_supported, Context, timeit -from tinygrad.lazy import LazyBuffer, view_supported_devices +from tinygrad.engine.lazy import LazyBuffer, view_supported_devices from extra.models.llama import precompute_freqs_cis class KernelCountException(Exception): pass diff --git a/test/test_subbuffer.py b/test/test_subbuffer.py index 8c0aaf3911..18702d53d4 100644 --- a/test/test_subbuffer.py +++ b/test/test_subbuffer.py @@ -1,7 +1,7 @@ import unittest from tinygrad import Device, dtypes, Tensor from tinygrad.device import Buffer -from tinygrad.lazy import view_supported_devices +from tinygrad.engine.lazy import view_supported_devices @unittest.skipIf(Device.DEFAULT not in view_supported_devices, "subbuffer not supported") class TestSubBuffer(unittest.TestCase): diff --git a/tinygrad/engine/graph.py b/tinygrad/engine/graph.py index 5b0cda1069..e8c5c28383 100644 --- a/tinygrad/engine/graph.py +++ b/tinygrad/engine/graph.py @@ -4,7 +4,7 @@ from typing import List, Any, DefaultDict from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MetaOps, TernaryOps, UOps, UOp from tinygrad.device import Device from tinygrad.helpers import GRAPHPATH, DEBUG, GlobalCounters -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer with contextlib.suppress(ImportError): import networkx as nx diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 4f544e3104..b9b928e644 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import TypeVar, Generic, Callable, List, Tuple, Union, Dict, cast, Optional, Any import functools, itertools, collections from tinygrad.tensor import Tensor -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, GRAPH, BEAM, getenv, colored, JIT, dedup, partition from tinygrad.device import Buffer, Compiled, Device from tinygrad.dtype import DType diff --git a/tinygrad/lazy.py b/tinygrad/engine/lazy.py similarity index 100% rename from tinygrad/lazy.py rename to tinygrad/engine/lazy.py diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index 9e73a15c17..9c39c9a427 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -7,7 +7,7 @@ from tinygrad.helpers import GRAPH, DEBUG, MULTIOUTPUT, SAVE_SCHEDULE, FUSE_CONV GlobalCounters, all_same, colored, prod, dedup, all_int, merge_dicts, getenv, Metadata, unwrap from tinygrad.shape.symbolic import Variable, sint from tinygrad.dtype import ImageDType, dtypes -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.device import Buffer from tinygrad.shape.view import View, strides_for_shape diff --git a/tinygrad/function.py b/tinygrad/function.py index 85761333c1..e4540eae60 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -5,7 +5,7 @@ from tinygrad.helpers import argsort from tinygrad.dtype import dtypes, DType, sum_acc_dtype from tinygrad.ops import ReduceOps, resolve from tinygrad.tensor import Function -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.shape.symbolic import sint class Contiguous(Function): diff --git a/tinygrad/multi.py b/tinygrad/multi.py index 5ccd096e2b..2ef0e63c8e 100644 --- a/tinygrad/multi.py +++ b/tinygrad/multi.py @@ -4,7 +4,7 @@ import functools, itertools, operator from tinygrad.helpers import all_same, all_int, dedup, prod, DEBUG, RING, getenv from tinygrad.dtype import DType from tinygrad.ops import REDUCE_ALU, BinaryOps, MetaOps, UnaryOps, TernaryOps, ReduceOps, MathTrait -from tinygrad.lazy import LazyBuffer +from tinygrad.engine.lazy import LazyBuffer from tinygrad.shape.shapetracker import sint def all_reduce(op: ReduceOps, lbs: List[LazyBuffer]) -> List[LazyBuffer]: diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index bd68b6cca1..e124626319 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -8,11 +8,11 @@ from collections import defaultdict from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype from tinygrad.helpers import argfix, make_pair, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup from tinygrad.helpers import IMAGE, DEBUG, WINO, _METADATA, Metadata, TRACEMETA -from tinygrad.lazy import LazyBuffer from tinygrad.multi import MultiLazyBuffer from tinygrad.ops import MetaOps, truncate, smax, resolve, UOp, UOps, BinaryOps from tinygrad.device import Device, Buffer, BufferOptions from tinygrad.shape.symbolic import sint, Variable +from tinygrad.engine.lazy import LazyBuffer from tinygrad.engine.realize import run_schedule, memory_planner from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars