github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

patch merge_views (#3311)

Browse Source
This commit is contained in:
Jyotirmaya Mahanta
2024-02-12 02:53:55 -08:00
committed by GitHub
parent b6a2600c86
commit d55f99e881
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tinygrad/shape/shapetracker.py
View File

@@ -18,8 +18,9 @@ def un1d(shape:Tuple[sint, ...], offs:sint) -> List[sint]:
@functools.lru_cache(maxsize=None)
def merge_views(vm2:View, vm1:View) -> Optional[View]:
if vm1.contiguous and vm1.shape == vm2.shape: return vm2
if vm2.contiguous: return vm1
if vm1.contiguous and vm1.shape == vm2.shape: return vm2
if vm1.contiguous and vm1.size() == vm2.size() and (ret := vm2.reshape(vm1.shape)) is not None: return ret
if not vm2.mask and vm1.offset == 0 and None not in (rstrides := ShapeTracker((vm2, vm1)).real_strides()):
return View.create(vm1.shape, cast(Tuple[sint, ...], rstrides), vm2.offset, vm1.mask)
if vm1.mask: