tinygrad 0.10.0__py3-none-any.whl → 0.10.1__py3-none-any.whl

{tinygrad-0.10.0.dist-info → tinygrad-0.10.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: tinygrad
-Version: 0.10.0
+Version: 0.10.1
 Summary: You like pytorch? You like micrograd? You love tinygrad! <3
 Author: George Hotz
 License: MIT
@@ -11,26 +11,19 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Provides-Extra: arm
 Requires-Dist: unicorn; extra == "arm"
-Provides-Extra: docs
-Requires-Dist: mkdocs; extra == "docs"
-Requires-Dist: mkdocs-material; extra == "docs"
-Requires-Dist: mkdocstrings[python]; extra == "docs"
-Requires-Dist: markdown-callouts; extra == "docs"
-Requires-Dist: markdown-exec[ansi]; extra == "docs"
-Requires-Dist: black; extra == "docs"
-Requires-Dist: numpy; extra == "docs"
+Provides-Extra: triton
+Requires-Dist: triton-nightly>=2.1.0.dev20231014192330; extra == "triton"
 Provides-Extra: linting
 Requires-Dist: pylint; extra == "linting"
-Requires-Dist: mypy==1.11.2; extra == "linting"
+Requires-Dist: mypy==1.13.0; extra == "linting"
 Requires-Dist: typing-extensions; extra == "linting"
 Requires-Dist: pre-commit; extra == "linting"
 Requires-Dist: ruff; extra == "linting"
 Requires-Dist: types-tqdm; extra == "linting"
-Provides-Extra: llvm
-Requires-Dist: llvmlite; extra == "llvm"
 Provides-Extra: testing
 Requires-Dist: numpy; extra == "testing"
 Requires-Dist: torch; extra == "testing"
+Requires-Dist: jax; extra == "testing"
 Requires-Dist: pillow; extra == "testing"
 Requires-Dist: pytest; extra == "testing"
 Requires-Dist: pytest-xdist; extra == "testing"
@@ -50,11 +43,28 @@ Requires-Dist: hypothesis; extra == "testing"
 Requires-Dist: nibabel; extra == "testing"
 Requires-Dist: bottle; extra == "testing"
 Requires-Dist: ggml-python; extra == "testing"
-Provides-Extra: testing_tf
+Requires-Dist: capstone; extra == "testing"
+Provides-Extra: webgpu
+Requires-Dist: wgpu; extra == "webgpu"
+Provides-Extra: docs
+Requires-Dist: mkdocs; extra == "docs"
+Requires-Dist: mkdocs-material; extra == "docs"
+Requires-Dist: mkdocstrings[python]; extra == "docs"
+Requires-Dist: markdown-callouts; extra == "docs"
+Requires-Dist: markdown-exec[ansi]; extra == "docs"
+Requires-Dist: black; extra == "docs"
+Requires-Dist: numpy; extra == "docs"
+Provides-Extra: testing-tf
 Requires-Dist: tensorflow==2.15.1; extra == "testing-tf"
-Requires-Dist: tensorflow-addons; extra == "testing-tf"
-Provides-Extra: triton
-Requires-Dist: triton-nightly>=2.1.0.dev20231014192330; extra == "triton"
+Requires-Dist: tensorflow_addons; extra == "testing-tf"
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: license
+Dynamic: provides-extra
+Dynamic: requires-python
+Dynamic: summary
 
 <div align="center">
 
@@ -146,6 +156,7 @@ tinygrad already supports numerous accelerators, including:
 - [x] [AMD](tinygrad/runtime/ops_amd.py)
 - [x] [NV](tinygrad/runtime/ops_nv.py)
 - [x] [QCOM](tinygrad/runtime/ops_qcom.py)
+- [x] [WEBGPU](tinygrad/runtime/ops_webgpu.py)
 
 And it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops.
 
@@ -183,8 +194,8 @@ y = Tensor([[2.0,0,-2.0]], requires_grad=True)
 z = y.matmul(x).sum()
 z.backward()
 
-print(x.grad.numpy())  # dz/dx
-print(y.grad.numpy())  # dz/dy
+print(x.grad.tolist())  # dz/dx
+print(y.grad.tolist())  # dz/dy
 ```
 
 The same thing but in PyTorch:
@@ -196,8 +207,8 @@ y = torch.tensor([[2.0,0,-2.0]], requires_grad=True)
 z = y.matmul(x).sum()
 z.backward()
 
-print(x.grad.numpy())  # dz/dx
-print(y.grad.numpy())  # dz/dy
+print(x.grad.tolist())  # dz/dx
+print(y.grad.tolist())  # dz/dy
 ```
 
 ## Contributing

tinygrad-0.10.1.dist-info/RECORD

@@ -0,0 +1,86 @@
+tinygrad/__init__.py,sha256=2Jhg7NSWlegCi4OAfGW0iHRVHeqMx09f7446rwAmc60,587
+tinygrad/device.py,sha256=mUrxoZJfqBJepDeqbmS2Y8UbOB_UPERqP9zBN7sCBk8,18968
+tinygrad/dtype.py,sha256=010zGuqXwUoyrhe23nJQ5oHmJt_6CRL7Hcl0kWjvKbo,9843
+tinygrad/gradient.py,sha256=hVyRMnwzjtWwKfF0NMGQovE2I7_2GdKFyR9pWJn4eE4,4280
+tinygrad/helpers.py,sha256=DIofGx-mg-umNKjhNpINcrxZExUwo6JIKHgmFuKNLUM,19203
+tinygrad/ops.py,sha256=T1gmR1ywWPnGEcw-MPAW5SGEvNS_789PtnFrR8rz5H4,70614
+tinygrad/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/spec.py,sha256=Wi8PdgeIHzv6WCTwKT9zGP5_53Bt_QtmvGTjIzTuwi0,8953
+tinygrad/tensor.py,sha256=-DVucOwq8HmM7IHqW7ZNABy7-vYZI-YfF0Vzdo0rs70,181559
+tinygrad/codegen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/codegen/kernel.py,sha256=4TcuMvvGODd0zFBuhK9KzvQhwzknRgeCL_n9WrOpljQ,42432
+tinygrad/codegen/linearize.py,sha256=X9aOMhvimDTvMq2Vnj3zfGNIRhJnAG5mLcq6EeUfvvU,10382
+tinygrad/codegen/lowerer.py,sha256=DqFg53oALqWG7RygOjXEioyGk52BnGbCzvQSqiOIKGw,7527
+tinygrad/codegen/rewriter.py,sha256=T3tOn-T1IFqiP90GCqfmkr0V5i0knr2tjQfKPVr6zzM,29948
+tinygrad/codegen/transcendental.py,sha256=0qRaEtIoJKDfjPqvQWexShZW3F__wtzfjhU__BqiMD8,13112
+tinygrad/engine/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/engine/jit.py,sha256=6D8pCpRPZNANwFEHAsAMUgUaEwaCf4mB_1oem6oZZ7o,16630
+tinygrad/engine/memory.py,sha256=UyiNYIoUjtcJ1SX6ApoCnbrSKZpbBbhfwr13TIaiqEM,3231
+tinygrad/engine/multi.py,sha256=FnyCvwFTdBXiydXB0dD5Y63KnerCurZS_o8WOU5fHFM,10371
+tinygrad/engine/realize.py,sha256=vOQ33sTn_P2oT0GRFU5tKYLfiRmxrIIdVs27v0_a40g,9690
+tinygrad/engine/schedule.py,sha256=9JywWlmVo912JmbIF9r5ajr5-Sesqw1EhE-CYnk5LHI,27407
+tinygrad/engine/search.py,sha256=VOsUOPaI7zWtDgL2SJDiKSPaTs3MhjgZOJBZ5_ISeiY,12093
+tinygrad/nn/__init__.py,sha256=BAxMz-g7v-v1A32KaBzmGiaEnvQ_dU9d3KoPdYYwLDQ,15156
+tinygrad/nn/datasets.py,sha256=wcT0Qrlpw_RzM7uBy8uphzKAjGT-ZE48fiP-0g3WvI0,1042
+tinygrad/nn/optim.py,sha256=qfdYKi_ssX5O_DU6h8GJ0WCzBzAZLyyS3p_946PJNsQ,6816
+tinygrad/nn/state.py,sha256=zXFMwAw7sf35C9x3iAxYiPsmhk7_S6qPcX3wXxqp6Bw,16030
+tinygrad/renderer/__init__.py,sha256=KvZ3y7MnqKDHOKoTcxghT95FhTBJuzs3IFEBM077jw8,6989
+tinygrad/renderer/cstyle.py,sha256=TrmxtioR7Lgw_8CK-Ay9EA0HjS1TXRf-RaMsCsuHszw,29852
+tinygrad/renderer/llvmir.py,sha256=iHFjE9-GtH2sGD_feCGe7aGECPHJ5qGOyuuyCDxSEXU,8529
+tinygrad/renderer/ptx.py,sha256=is4PMkdbDDgKsTy67DtkugatdhbtVMijKE9u1f6-0ag,14899
+tinygrad/renderer/wgsl.py,sha256=3AGZp2jvymxLZSJHjzD7zlwxvds36-lSu2ZkntrW_ww,6697
+tinygrad/runtime/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/runtime/ops_amd.py,sha256=zRictmcgUPw6d8iykz6QWqCmiPirOh-hTACT-GYxso8,37282
+tinygrad/runtime/ops_clang.py,sha256=InJPVVhY10bxRXOaDQRUBlWafSZRdo72Y0VVW9O6dBw,1310
+tinygrad/runtime/ops_cloud.py,sha256=1umLappzYKg0ECC9C_EN1lzXquAQVqY7wj-mS0vTcSk,10362
+tinygrad/runtime/ops_cuda.py,sha256=26ko2bTLT_jDmU2i8_xiL4BomD9krzvlz6wY5FE-e5c,7138
+tinygrad/runtime/ops_disk.py,sha256=3eTpJSVQQH32KBa7tRLBm2pPSkCaIC1adPc748UkUeg,6619
+tinygrad/runtime/ops_dsp.py,sha256=JCSehSDLzaa8Nhd6Xtas0d5vOCgwQGY4U3awHOLThEI,16900
+tinygrad/runtime/ops_gpu.py,sha256=VrY9iM5i44OQwTJE8sgUavCtr1TOpE96I5-JmqUFz-E,8972
+tinygrad/runtime/ops_hip.py,sha256=MbR4depxgHcaGpOJGvUCiLq5tdbpaDiqs-Xj41rY2xQ,3730
+tinygrad/runtime/ops_llvm.py,sha256=yjW5wVMJdmayeeTQwnQ60mtOWiiYEqvtb9Rt5cDj2rU,3239
+tinygrad/runtime/ops_metal.py,sha256=rz3fiaiB1by_RuP7VFVBgGA1mye43_-6PIfMfeQ3YmU,13229
+tinygrad/runtime/ops_npy.py,sha256=8VNf1S5M_MRk9d3GxSsTPbfEz7I_aOwl7QMZ1mUG3As,370
+tinygrad/runtime/ops_nv.py,sha256=8NvxtEc89dbJ5xgVhGs9_zeGMztp9HRZHLwPqylDsdQ,34541
+tinygrad/runtime/ops_python.py,sha256=5CZek-7fMhq6w07lGDM0pXltZzzMMRtuU2x5PeWMyWY,11681
+tinygrad/runtime/ops_qcom.py,sha256=Dt4hgAd6o13CxsOFRSt7lHY6bCOTLvtQpOr_jx_lYbc,22565
+tinygrad/runtime/ops_webgpu.py,sha256=-uBAWhVZ7T_9zG2v6PWjkJLpovrClqb_XAPk-l83ryc,4345
+tinygrad/runtime/autogen/adreno.py,sha256=u7VxIomPAlW3nFUs4gSTe-6ijam_ywkvDM9OuTLF-j8,897915
+tinygrad/runtime/autogen/amd_gpu.py,sha256=Iasq-zYiv8bvT43dtvPO1W5jaLEQ3d6hP0CoFVhSsak,3977783
+tinygrad/runtime/autogen/comgr.py,sha256=3pp3XyqEJDBLa9XtGx2-Gc1iJgBbbgIq4pdFEpYXT44,39874
+tinygrad/runtime/autogen/cuda.py,sha256=N0QyaMvQumr_HZh7fusCHM1d4o4mYti3Wq1MN7JSKr8,243920
+tinygrad/runtime/autogen/hip.py,sha256=1yUHDCwL3KkD15if2Q1Ud3GbJiR7DxsNorKZTCINw54,245532
+tinygrad/runtime/autogen/hsa.py,sha256=7Hsrn17HmChyeFOSX_3Fnzl9c0COtq2Z2ExqGu5FNiU,277716
+tinygrad/runtime/autogen/io_uring.py,sha256=ZIZ2YnQkLr8WIHMieBw9Dv-NZ1ar9TwdP4YBv3gJm28,59786
+tinygrad/runtime/autogen/kfd.py,sha256=VdhuG4qec0EgM-jJmWcdTS-8WrmywNkcjSX7ibbmvdk,30866
+tinygrad/runtime/autogen/kgsl.py,sha256=2EgJ5Kst4oRUv81hsV2srgwPvWpY-weaSB4E2lGMAyc,50656
+tinygrad/runtime/autogen/libc.py,sha256=xKJk2hCzVpauJSc8wCQis5x3SwcXnDli7_HyRUqEGRc,197318
+tinygrad/runtime/autogen/libpciaccess.py,sha256=zaVmkoUHVTEQcPQwkFpgMCgX4HtX-BL-MGMbi8XtgCI,84194
+tinygrad/runtime/autogen/llvm.py,sha256=aeVd_ByohxbGRyqXzShPOupI2xtcdk34I6_OIBrMQHg,467606
+tinygrad/runtime/autogen/nv_gpu.py,sha256=9X2tPdv2E5JmXGZeT8i9jL19YJ4ETTsYwfU_Wn9mTwc,1679326
+tinygrad/runtime/autogen/nvrtc.py,sha256=19te2-TW5suFy85KnJox3CPOmeeml5YxqIDeL-Bx_m4,23132
+tinygrad/runtime/autogen/opencl.py,sha256=NL6fa8P3KC_McNZ8g2babdr3b8vrY-bFK0qzNAtL-rE,82656
+tinygrad/runtime/autogen/qcom_dsp.py,sha256=jx36-zC6reTuWgfbHCrKVjOZcF4Q9fBnq3CuTbxztQk,61848
+tinygrad/runtime/autogen/vfio.py,sha256=IJV1eeWWllU6b9LAX_IH0bUW5NDzfhPQy_YzXGhD9-8,32431
+tinygrad/runtime/graph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/runtime/graph/cuda.py,sha256=vLjT_c93G6ia_1MsbYWP5Uq96Aeko0AOskRkwT5-MUI,4818
+tinygrad/runtime/graph/hcq.py,sha256=kKu2YnjAZU40XMACSbbcxJSi2xdTg3OYLO2zcPLyAf0,12600
+tinygrad/runtime/graph/metal.py,sha256=6JN7WJr9w1pIvQGYr0Rsnyg-SA-slK61jKV__KmHnXg,6001
+tinygrad/runtime/support/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/runtime/support/allocator.py,sha256=INW6TaJxMi4cwMEDtkta2wFH3kJ87wy3SyFFu3_jJ9w,4721
+tinygrad/runtime/support/compiler_cuda.py,sha256=6cU1OMMW3aOUFNVALpDYWKXh-zFc5q81PSSQhRK9fLw,5471
+tinygrad/runtime/support/compiler_hip.py,sha256=fbRP82UdG4T-KCRYH_H2hEXlMFeHIJntSnY35ZWE5JY,4398
+tinygrad/runtime/support/elf.py,sha256=AxWyaAVEe4xdSTiISqIf80oaHRwUZwb5b1V_Q3874s8,3857
+tinygrad/runtime/support/hcq.py,sha256=JWLRwoGPZCgRBryNb6uq0qCcfkVfMmNnlnHmCzPFxFI,21901
+tinygrad/runtime/support/llvm.py,sha256=-qI8NRmSx2dBZnG-OQCTWfmy0XySj_T8kHNRHULTyG4,2174
+tinygrad/runtime/support/am/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/runtime/support/am/amdev.py,sha256=LD80Bc9DHdl99bxrfwn2c43m1NdDD0mjPnSOJkEfkgU,20433
+tinygrad/runtime/support/am/ip.py,sha256=WnSIZWSG9IvSjVYLJv-VShW_X6vcziynj__lvTOt4yQ,24531
+tinygrad/shape/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/shape/shapetracker.py,sha256=clrYjN-zEZHSh_biDy3jnHL4Fq9HyWCdH_CRpwLKki0,7612
+tinygrad/shape/view.py,sha256=7KJwP2lS1YWhMq80Ka6_h_qavyxSEm9jWloVgHYRx-k,18110
+tinygrad-0.10.1.dist-info/LICENSE,sha256=ABRhUPEILzINYIukgazD-_rPipkUNUwslrb0RxnV6Xc,1058
+tinygrad-0.10.1.dist-info/METADATA,sha256=JMb7PpBcZqsf0tbmNJzvBqUrH6U1JER_0mz20UMo0LA,11241
+tinygrad-0.10.1.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+tinygrad-0.10.1.dist-info/top_level.txt,sha256=vDABMCWBFQnx2kn9Azueu88FP-1klQdePoHikQhHymc,9
+tinygrad-0.10.1.dist-info/RECORD,,

{tinygrad-0.10.0.dist-info → tinygrad-0.10.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.5.0)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tinygrad/engine/lazy.py

@@ -1,228 +0,0 @@
-from __future__ import annotations
-from typing import Optional, Any, Tuple, List, get_args
-from tinygrad.dtype import dtypes, DType, ConstType, to_dtype, ImageDType
-from tinygrad.helpers import prod, getenv, all_int, all_same, DEBUG, _METADATA, Metadata, SPLIT_REDUCEOP, LAZYCACHE
-from tinygrad.ops import exec_alu, python_alu
-from tinygrad.ops import identity_element, MathTrait, resolve, UOp, sint, GroupOp, Ops
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.device import Buffer
-from weakref import ref, ReferenceType, WeakValueDictionary
-
-lazycache: WeakValueDictionary[Any, LazyBuffer] = WeakValueDictionary()
-def create_lazybuffer(device:str, st:ShapeTracker, dtype:DType, op:Optional[Ops]=None, arg:Any=None, srcs:Tuple[LazyBuffer, ...]=(),
-                      base:Optional[LazyBuffer]=None, enable_cache=bool(LAZYCACHE)):
-  if st.size == 0: op, arg, srcs, base = Ops.CONST, 0, (), None
-  dtype = to_dtype(dtype)
-  if op is Ops.CONST: arg, enable_cache = dtypes.as_const(arg, dtype) if not isinstance(arg, UOp) else arg, True
-
-  cache_key = (device, st, dtype, op, arg, tuple(ref(x) for x in srcs)) if base is None else (st, ref(base))
-  if enable_cache and (rret := lazycache.get(cache_key, None)) is not None: return rret
-
-  ret = LazyBuffer(device, st, dtype, op, arg, srcs, base=base, metadata=_METADATA.get())
-  if enable_cache: lazycache[cache_key] = ret
-  return ret
-
-view_supported_devices = {"LLVM", "CLANG", "CUDA", "NV", "AMD", "METAL", "QCOM", "DSP", "DISK"}
-class LazyBuffer(MathTrait):
-  def __init__(self, device:str, st:ShapeTracker, dtype:DType,
-               op:Optional[Ops]=None, arg:Any=None, srcs:Tuple[LazyBuffer, ...]=(),
-               base:Optional[LazyBuffer]=None, metadata:Optional[Metadata]=None):
-    self.device, self.st, self.dtype, self.shape, self.size, self.metadata = device, st, to_dtype(dtype), st.shape, st.size, metadata
-    self._base: Optional[LazyBuffer] = None
-    if base is None:
-      # properties on base
-      self.op, self.arg, self.srcs = op, arg, srcs  # this is a UOp, except the src is LazyBuffers and not UOps
-      assert self.op is not Ops.ASSIGN or srcs[0].base.realized is not None, "assign target must be realized"
-
-      if self.op is Ops.BUFFER_VIEW:
-        # some LazyBuffers can be processed with only a view, no AST required
-        self.buffer: Buffer = srcs[0].base.buffer.view(st.size, self.dtype, srcs[0].st.views[0].offset * srcs[0].dtype.itemsize)
-      else:
-        self.buffer = srcs[0].base.buffer if self.op is Ops.ASSIGN else Buffer(device, self.size, self.dtype)
-      self.buffer.ref(1)
-      self.contiguous_child: Optional[Tuple[ReferenceType[LazyBuffer], ShapeTracker]] = None
-      self.forced_realize = False
-    else:
-      # properties on view
-      assert base.base == base, "base must be a base itself"
-      self._base = base
-
-  def __del__(self):
-    if hasattr(self, 'buffer'): self.buffer.ref(-1)
-
-  def __repr__(self) -> str:
-    return f"<LB {self.device} {self.shape} {str(self.dtype)[7:]} {self.st if self.base is not self else (self.op, self.realized)}>"
-
-  @property
-  def realized(self) -> Optional[Buffer]:
-    # NOTE: we check for a lack of srcs instead of an allocated buffer to make unrealized assigns return None here
-    return self.buffer if self._base is None and not hasattr(self, 'srcs') else None
-
-  # NOTE: this has to be a function to prevent self reference
-  @property
-  def base(self) -> LazyBuffer: return self._base if self._base is not None else self
-
-  # same API as multi
-  @property
-  def lbs(self) -> List[LazyBuffer]: return [self]
-
-  @staticmethod
-  def metaop(op, shape:Tuple[sint,...], dtype:DType, device:str, arg=None, src:Tuple[LazyBuffer, ...]=(), enable_cache=False) -> LazyBuffer:
-    assert isinstance(src, tuple)
-    return create_lazybuffer(device, ShapeTracker.from_shape(shape), dtype, op, arg, src, enable_cache=enable_cache)
-
-  def const_like(self, b): return self.const_with_shape(b, self.shape)
-  def const_with_shape(self, val:ConstType, shape:Tuple[sint,...]) -> LazyBuffer:
-    assert isinstance(val, get_args(ConstType)), f"{val=} has {type(val)=}, not a ConstType"
-    return LazyBuffer.metaop(Ops.CONST, tuple(), self.dtype, self.device, arg=val).reshape((1,)*len(shape)).expand(shape)
-
-  @property
-  def is_realized(self) -> bool: return self.base.realized is not None
-
-  def assign(self, x:LazyBuffer) -> LazyBuffer:
-    assert x.size == self.size, f"assign target must have same size {self.size=} != {x.size=}"
-    assert self.is_realized, f"assign target must be realized {self}"
-    return LazyBuffer.metaop(Ops.ASSIGN, self.shape, self.dtype, self.device, arg=() if self.st.contiguous else (self.st,),
-                             src=(self.base, x), enable_cache=True)
-
-  def can_view(self):
-    return (self.st.consecutive and not self.is_unrealized_const() and not isinstance(self.dtype, ImageDType) and
-            self.device.split(":")[0] in view_supported_devices)
-
-  def contiguous(self, allow_buffer_view=True):
-    if not self.st.contiguous or self.size != self.base.size or self.is_unrealized_const():
-      ret = self.alu(Ops.BUFFER_VIEW) if allow_buffer_view and self.can_view() else self.alu(Ops.CONTIGUOUS)
-      if (sti := self.st.invert(self.base.shape)) is not None: self.base.contiguous_child = ref(ret), sti
-      return ret
-    self.base.forced_realize = True
-    return self
-
-  def bitcast(self, dtype:DType) -> LazyBuffer: return self.cast(dtype, bitcast=True)
-  def cast(self, dtype:DType, bitcast:bool=False, allow_buffer_view=True) -> LazyBuffer:
-    if self.dtype == dtype: return self
-    if self.device.startswith("DISK") and not bitcast: raise RuntimeError("attempted to cast disk buffer (bitcast only)")
-    if self.is_unrealized_unmasked_const() and not bitcast:
-      return create_lazybuffer(self.device, self.st, dtype, Ops.CONST, dtypes.as_const(self.base.arg, dtype))
-    new_shape = self.shape
-    if bitcast and self.dtype.itemsize != dtype.itemsize:
-      if not self.device.startswith("DISK"): raise RuntimeError("shape changing bitcast only supported on DISK right now")
-      if not all_int(new_shape): raise RuntimeError("shape changing bitcast with symbolic shape isn't supported yet")
-      # https://pytorch.org/docs/stable/generated/torch.Tensor.view.html
-      if not (new_shape[-1]*self.dtype.itemsize) % dtype.itemsize == 0: raise RuntimeError("unsupported size in bitcast")
-      new_shape = new_shape[:-1] + ((new_shape[-1]*self.dtype.itemsize) // dtype.itemsize,)
-    elif getenv("CAST_BEFORE_VIEW", 1) and dtype.itemsize <= self.dtype.itemsize and self is not self.base:
-      # TODO: applying this makes gpt2 slower
-      return self.base.cast(dtype, bitcast)._view(self.st)
-    cast_op: Ops = (Ops.BUFFER_VIEW if self.can_view() and allow_buffer_view else Ops.BITCAST) if bitcast else Ops.CAST
-    return create_lazybuffer(self.device, ShapeTracker.from_shape(new_shape), dtype, cast_op, dtype, (self,))
-
-  def is_unrealized_const(self): return self.base.realized is None and self.base.op is Ops.CONST and not isinstance(self.base.arg, UOp)
-  def is_unrealized_unmasked_const(self): return self.is_unrealized_const() and all(v.mask is None for v in self.st.views)
-
-  def _copy(self, device:str) -> LazyBuffer:
-    assert self.st.contiguous and self.size == self.base.size, f"can only copy contig {self} {self.base}"
-    return create_lazybuffer(device, ShapeTracker.from_shape(self.shape), self.dtype, Ops.COPY, self.buffer.nbytes, (self,), enable_cache=False)
-
-  def copy_to_device(self, device:str, force:bool=False, clone:bool=False) -> LazyBuffer:
-    # no COPY
-    if self.device == device and not clone: return self
-
-    # double COPY = one COPY
-    if not force and self.st.contiguous and self.size == self.base.size and not self.base.realized and self.base.op is Ops.COPY:
-      return self.base.srcs[0].copy_to_device(device).reshape(self.st.shape)
-
-    # const doesn't have to be copied (issues with disk tensor)
-    if self.is_unrealized_const():
-      return LazyBuffer.metaop(Ops.CONST, tuple(), self.dtype, device, arg=self.base.arg)._view(self.st)
-
-    # if it's a shrink, do the shrink before the copy with CONTIGUOUS
-    if prod(self.st.shape) < prod(self.base.st.shape): return self.contiguous()._copy(device)
-
-    # copy the base and apply the shapetracker on the new device
-    return self.base._copy(device)._view(self.st)
-
-  def clone(self) -> LazyBuffer: return self.copy_to_device(self.device, clone=True)
-
-  def alu(self, op:Ops, *in_srcs:LazyBuffer) -> LazyBuffer:
-    srcs: List[LazyBuffer] = []
-    for s in (self,)+in_srcs:
-      if s == s.base and s.base.contiguous_child and (root:=s.base.contiguous_child[0]()) is not None:
-        srcs.append(root._view(s.base.contiguous_child[1]))
-      else:
-        srcs.append(s)
-    if not all_same(dts:=[x.dtype.base for x in (srcs[1:] if op is Ops.WHERE else srcs)]):
-      raise AssertionError(f"all dtypes must match {dts} on {op}")
-    assert all_same([x.shape for x in srcs]), f"all shapes must be the same {[x.shape for x in srcs]}"
-    if op is Ops.WHERE: assert srcs[0].dtype == dtypes.bool, "Ops.WHERE must have the first arg be bool"
-
-    out_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else srcs[-1].dtype
-
-    # const folding
-    if op in python_alu and all(s.is_unrealized_unmasked_const() for s in srcs):
-      return self.cast(out_dtype).const_like(exec_alu(op, out_dtype, [s.base.arg for s in srcs]))
-    if op in GroupOp.Binary:
-      x, y = self, in_srcs[0]
-      if op is Ops.ADD:
-        if y.is_unrealized_unmasked_const() and y.base.arg == 0: return x
-        if x.is_unrealized_unmasked_const() and x.base.arg == 0: return y
-      if op is Ops.MUL:
-        if x.is_unrealized_unmasked_const() and (val := x.base.arg) in (1, 0): return y if val == 1 else y.const_like(0)
-        if y.is_unrealized_unmasked_const() and (val := y.base.arg) in (1, 0): return x if val == 1 else x.const_like(0)
-      if op is Ops.IDIV and y.is_unrealized_unmasked_const() and y.base.arg == 1: return x
-
-    return create_lazybuffer(self.device, ShapeTracker.from_shape(self.shape), out_dtype, op, None, tuple(srcs))
-
-  # *** reduce ops ***
-
-  def _reduce_op(self, op:Ops, axis:Tuple[int, ...]) -> LazyBuffer:
-    assert all(0 <= x < len(self.shape) for x in axis), f"axis args {axis} out of range for shape {self.shape}"
-    axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
-    if len(axis) == 0: return self
-    return create_lazybuffer(self.device, ShapeTracker.from_shape(self.st.reduce(axis)), self.dtype, Ops.REDUCE_AXIS, (op, axis), (self,))
-
-  def r(self, op:Ops, axis:Tuple[int, ...]) -> LazyBuffer:
-    new_shape = self.st.reduce(axis)
-    # TODO: this logic should move to the scheduler
-    if 0 in self.shape and 0 not in new_shape: return self.const_with_shape(identity_element(op, self.dtype), new_shape)
-
-    # const folding
-    # TODO: fold this for symbolic?
-    if self.is_unrealized_unmasked_const() and all_int(self.shape):
-      if op is Ops.ADD: return self.const_with_shape(self.base.arg * prod(self.shape[i] for i in axis), new_shape)
-      if op is Ops.MUL: return self.const_with_shape(self.base.arg ** prod(self.shape[i] for i in axis), new_shape)
-      if op is Ops.MAX: return self.const_with_shape(self.base.arg, new_shape)
-
-    # TODO: can we split symbolic shape if the reduce axis is not symbolic?
-    if not SPLIT_REDUCEOP or not all_int(self.shape) or (0 in self.shape) or \
-      prod(self.shape) // prod(new_shape) < getenv("REDUCEOP_SPLIT_THRESHOLD", 32768):
-      return self._reduce_op(op, axis)
-
-    # if there are few globals, make some reduces into globals by splitting into two kernels
-    # cap output buffer to 2**22: heuristic number of global outputs to achieve max occupancy with enough locals+upcasts for gemm
-    #   ~2**10 should be enough if GROUP is used
-    # 256 split maximum should be "negligible reduce" for low prod(new_shape), 8 split minimum.
-    # split is moved to the end to provide maximum locality for the second phase reduce.
-    self_real_strides = self.st.real_strides(ignore_valid=True)
-    split_candidates = [(i, x) for i in axis for x in range(min(256,2**getenv("REDUCEOP_SPLIT_SIZE",22)//prod(new_shape)),8-1,-1)
-                        if self.shape[i] % x == 0 and self_real_strides[i] != 0]
-    if not split_candidates: return self._reduce_op(op, axis)
-    dim_to_split, divisor = split_candidates[0]
-    splitted_shape = self.shape[:dim_to_split] + (divisor,) + (self.shape[dim_to_split]//divisor,) + self.shape[dim_to_split+1:]
-    splitted = self.reshape(splitted_shape).permute(tuple([x for x in range(len(splitted_shape)) if x != dim_to_split]+[dim_to_split]))
-    if DEBUG >= 3: print(f"split {divisor}: {self.shape} -> {splitted.shape} -> {new_shape}")
-    return splitted._reduce_op(op, axis)._reduce_op(op, (len(new_shape),)).reshape(new_shape)  # reduce original axes, then split
-
-  # *** movement ops ***
-
-  def _view(self, new_st:ShapeTracker) -> LazyBuffer:
-    if self.st.size == 0 or (new_st.views[-1].mask is not None and any((x[1]-x[0]) == 0 for x in new_st.views[-1].mask)):
-      return self.const_with_shape(0, new_st.shape)
-    if new_st.contiguous and self.base.shape == new_st.shape: return self.base
-    return create_lazybuffer(self.device, new_st, self.dtype, base=self.base)
-
-  def reshape(self, arg:Tuple[sint, ...]): return self._view(self.st.reshape(arg))
-  def pad(self, arg:Tuple[Tuple[sint, sint], ...]): return self._view(self.st.pad(arg))
-  def expand(self, arg:Tuple[sint, ...]): return self._view(self.st.expand(arg))
-  def permute(self, arg:Tuple[int, ...]): return self._view(self.st.permute(arg))
-  def shrink(self, arg:Tuple[Tuple[sint, sint], ...]): return self._view(self.st.shrink(arg))
-  def stride(self, arg:Tuple[int, ...]): return self._view(self.st.stride(arg))

tinygrad/function.py

@@ -1,212 +0,0 @@
-"""This is where the forwards and backwards passes live."""
-import math
-from typing import Tuple, Optional
-from tinygrad.helpers import argsort
-from tinygrad.dtype import dtypes, DType, sum_acc_dtype
-from tinygrad.ops import Ops, resolve, sint
-from tinygrad.tensor import Function
-from tinygrad.engine.lazy import LazyBuffer
-
-class Contiguous(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer: return x.contiguous()
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output
-
-class ContiguousBackward(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer: return x
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.contiguous()
-
-class Cast(Function):
-  def forward(self, x:LazyBuffer, dtype:DType, bitcast:bool=False) -> LazyBuffer:
-    self.input_dtype, self.bitcast = x.dtype, bitcast
-    return x.bitcast(dtype) if self.bitcast else x.cast(dtype)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    if self.bitcast: raise RuntimeError("bitcast cannot backward")
-    return grad_output.cast(self.input_dtype)
-
-# ************* unary ops *************
-
-class Reciprocal(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.ret = x.reciprocal()
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return -grad_output * self.ret * self.ret
-
-class Sin(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.x = x
-    return x.sin()
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return (math.pi/2 - self.x).sin() * grad_output
-
-class Relu(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.ret = x.maximum(0)
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return self.ret.gt(0).cast(grad_output.dtype) * grad_output
-
-class Log(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.x = x
-    return x.log2() * math.log(2)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output / self.x
-
-class Exp(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.ret = (x * (1/math.log(2))).exp2()
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return self.ret * grad_output
-
-class Sqrt(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.ret = x.sqrt()
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output / (self.ret*2)
-
-# NOTE: the implicit derivative of sigmoid is not stable
-# https://towardsdatascience.com/derivative-of-the-sigmoid-function-536880cf918e
-# TODO: have the backend automatically find this
-class Sigmoid(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.ret = (1 + (x * (-1/math.log(2))).exp2()).reciprocal()
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    return (self.ret * (1 - self.ret)) * grad_output
-
-class Sign(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer: return x.ne(0).where(x.lt(0).where(x.const_like(-1), x.const_like(1)), x.const_like(0))
-  # backward always return 0 to match torch
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.const_like(0)
-
-# ************* binary ops *************
-
-class Less(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x.lt(y)
-  def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]: return None, None
-
-class Neq(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x.ne(y)
-  def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]: return None, None
-
-class Xor(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x^y
-
-class BitwiseAnd(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x&y
-
-class BitwiseOr(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x|y
-
-class Threefry(Function):
-  def forward(self, x:LazyBuffer, seed:LazyBuffer) -> LazyBuffer: return x.threefry(seed)
-
-class Add(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x+y
-
-  def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]:
-    return grad_output if self.needs_input_grad[0] else None, \
-           grad_output if self.needs_input_grad[1] else None
-
-class Mul(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer:
-    self.x, self.y = x, y
-    return x * y
-
-  def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]:
-    return (self.y * grad_output) if self.needs_input_grad[0] else None, \
-           (self.x * grad_output) if self.needs_input_grad[1] else None
-
-class IDiv(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x // y
-
-# ************* ternary ops *************
-
-class Where(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer, z:LazyBuffer) -> LazyBuffer:
-    self.x = x
-    return self.x.where(y, z)
-
-  def backward(self, grad_output:LazyBuffer) -> Tuple[None, Optional[LazyBuffer], Optional[LazyBuffer]]:
-    return None, \
-      self.x.where(grad_output, grad_output.const_like(0)) if self.needs_input_grad[1] else None, \
-      self.x.where(grad_output.const_like(0), grad_output) if self.needs_input_grad[2] else None
-
-# ************* reduce ops *************
-
-class Sum(Function):
-  def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer:
-    self.input_shape = x.shape
-    return x.r(Ops.ADD, axis)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.expand(self.input_shape)
-
-class Prod(Function):
-  def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer:
-    self.x, self.ret = x, x.r(Ops.MUL, axis)
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    return (grad_output * self.ret).expand(self.x.shape) / self.x
-
-class Max(Function):
-  def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer:
-    self.x, self.ret, self.axis = x, x.r(Ops.MAX, axis), axis
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    # 1s in locations where the max was chosen (can be two locations)
-    max_is_1s = self.x.ne(self.ret.expand(self.x.shape)).ne(self.x.const_like(1).cast(dtypes.bool)).cast(grad_output.dtype)
-    div = max_is_1s.r(Ops.ADD, self.axis).expand(self.x.shape)
-    return (max_is_1s/div) * grad_output.expand(self.x.shape)
-
-# ************* movement ops *************
-
-# NOTE: this is sum in reverse
-class Expand(Function):
-  def forward(self, x:LazyBuffer, shape:Tuple[int, ...]) -> LazyBuffer:
-    self.expanded_axis = tuple(i for i, (si, so) in enumerate(zip(x.shape, shape)) if resolve(si != so))
-    return x.expand(shape)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    return grad_output.cast(sum_acc_dtype(grad_output.dtype)).r(Ops.ADD, self.expanded_axis).cast(grad_output.dtype)
-
-class Reshape(Function):
-  def forward(self, x:LazyBuffer, shape:Tuple[int, ...]) -> LazyBuffer:
-    self.input_shape = x.shape
-    return x.reshape(shape)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.reshape(self.input_shape)
-
-class Permute(Function):
-  def forward(self, x:LazyBuffer, order:Tuple[int, ...]) -> LazyBuffer:
-    self.input_order = order
-    return x.permute(order)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.permute(argsort(self.input_order))
-
-class Pad(Function):
-  def forward(self, x:LazyBuffer, arg:Tuple[Tuple[int, int], ...]) -> LazyBuffer:
-    self.narg = tuple([(p[0], s+p[0]) for s,p in zip(x.shape, arg)])
-    return x.pad(arg)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.shrink(self.narg)
-
-class Shrink(Function):
-  def forward(self, x:LazyBuffer, arg:Tuple[Tuple[sint, sint], ...]) -> LazyBuffer:
-    self.narg = tuple([(p[0], s-p[1]) for s,p in zip(x.shape, arg)])
-    return x.shrink(arg)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.pad(self.narg)
-
-class Flip(Function):
-  def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer:
-    self.arg = tuple([-1 if i in axis else 1 for i in range(len(x.shape))])
-    return x.stride(self.arg)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.stride(self.arg)