fix: bump build

build/torch26-cxx11-cu118-x86_64-linux/megablocks/{_megablocks_dabb815.abi3.so → _megablocks_0586ba6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7a20cd4dc15095b8504db981c651e516e8a7d8394b99d973d632558637c8dba9
+oid sha256:7fbec6fa49d1b926d45b39b7e8393e06ee9622d0012501adaec213cb5802c86d
 size 10517576

build/torch26-cxx11-cu118-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_dabb815
-ops = torch.ops._megablocks_dabb815
+from . import _megablocks_0586ba6
+ops = torch.ops._megablocks_0586ba6
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_dabb815::{op_name}"
+    return f"_megablocks_0586ba6::{op_name}"

build/torch26-cxx11-cu118-x86_64-linux/megablocks/ops/binned_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu118-x86_64-linux/megablocks/ops/binned_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu118-x86_64-linux/megablocks/ops/gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu118-x86_64-linux/megablocks/ops/padded_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu118-x86_64-linux/megablocks/ops/padded_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu118-x86_64-linux/megablocks/ops/scatter.py

@@ -4,7 +4,7 @@
 from typing import Any, Optional
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu118-x86_64-linux/megablocks/ops/stk_autocast.py

@@ -0,0 +1,39 @@
+# vendored from
+# https://github.com/stanford-futuredata/stk/blob/736313768ef697ce13a0594a41b2512a0fbc9884/stk/backend/autocast.py
+import functools
+import torch
+
+
+def _is_eligible(x):
+    return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64)
+
+
+def _cast(x, dtype):
+    if isinstance(x, torch.Tensor) and _is_eligible(x):
+        return x.to(dtype)
+    elif isinstance(x, map):
+        return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()}
+    elif isinstance(x, list) or isinstance(x, tuple):
+        return type(x)(map(lambda y: _cast(y, dtype), x))
+    return x
+
+
+def custom_fwd(fwd):
+    """Wrap a custom autograd function that always uses autocast dtype."""
+
+    @functools.wraps(fwd)
+    def decorate_fwd(*args, **kwargs):
+        if torch.is_autocast_enabled():
+            with torch.autocast(device_type="cuda", enabled=False):
+                dtype = torch.get_autocast_gpu_dtype()
+                return fwd(*_cast(args, dtype), **_cast(kwargs, dtype))
+        return fwd(*args, **kwargs)
+    return decorate_fwd
+
+
+def custom_bwd(bwd):
+    @functools.wraps(bwd)
+    def decorate_bwd(*args, **kwargs):
+        with torch.autocast(device_type="cuda", enabled=False):
+            return bwd(*args, **kwargs)
+    return decorate_bwd

build/torch26-cxx11-cu124-x86_64-linux/megablocks/{_megablocks_dabb815.abi3.so → _megablocks_0586ba6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a3f69e5978b727f08b43112c2321a222719aa824612d452029225a48976dfbb6
+oid sha256:16141033c118b488348a29f3436f778764f8f4275fe510dc36badb7c152e0d42
 size 11869392

build/torch26-cxx11-cu124-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_dabb815
-ops = torch.ops._megablocks_dabb815
+from . import _megablocks_0586ba6
+ops = torch.ops._megablocks_0586ba6
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_dabb815::{op_name}"
+    return f"_megablocks_0586ba6::{op_name}"

build/torch26-cxx11-cu124-x86_64-linux/megablocks/ops/binned_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu124-x86_64-linux/megablocks/ops/binned_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu124-x86_64-linux/megablocks/ops/gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu124-x86_64-linux/megablocks/ops/padded_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu124-x86_64-linux/megablocks/ops/padded_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu124-x86_64-linux/megablocks/ops/scatter.py

@@ -4,7 +4,7 @@
 from typing import Any, Optional
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu124-x86_64-linux/megablocks/ops/stk_autocast.py

@@ -0,0 +1,39 @@
+# vendored from
+# https://github.com/stanford-futuredata/stk/blob/736313768ef697ce13a0594a41b2512a0fbc9884/stk/backend/autocast.py
+import functools
+import torch
+
+
+def _is_eligible(x):
+    return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64)
+
+
+def _cast(x, dtype):
+    if isinstance(x, torch.Tensor) and _is_eligible(x):
+        return x.to(dtype)
+    elif isinstance(x, map):
+        return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()}
+    elif isinstance(x, list) or isinstance(x, tuple):
+        return type(x)(map(lambda y: _cast(y, dtype), x))
+    return x
+
+
+def custom_fwd(fwd):
+    """Wrap a custom autograd function that always uses autocast dtype."""
+
+    @functools.wraps(fwd)
+    def decorate_fwd(*args, **kwargs):
+        if torch.is_autocast_enabled():
+            with torch.autocast(device_type="cuda", enabled=False):
+                dtype = torch.get_autocast_gpu_dtype()
+                return fwd(*_cast(args, dtype), **_cast(kwargs, dtype))
+        return fwd(*args, **kwargs)
+    return decorate_fwd
+
+
+def custom_bwd(bwd):
+    @functools.wraps(bwd)
+    def decorate_bwd(*args, **kwargs):
+        with torch.autocast(device_type="cuda", enabled=False):
+            return bwd(*args, **kwargs)
+    return decorate_bwd

build/torch26-cxx11-cu126-x86_64-linux/megablocks/{_megablocks_dabb815.abi3.so → _megablocks_0586ba6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:57540f7b6eae09c2c62826d13dfa2be53eaa37c86206df5914611a3fad9878ba
+oid sha256:3ea768d3d4780563159dd50075ed14d51166e5c3de9f5bd132047cfa6a23ef48
 size 11931048

build/torch26-cxx11-cu126-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_dabb815
-ops = torch.ops._megablocks_dabb815
+from . import _megablocks_0586ba6
+ops = torch.ops._megablocks_0586ba6
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_dabb815::{op_name}"
+    return f"_megablocks_0586ba6::{op_name}"

build/torch26-cxx11-cu126-x86_64-linux/megablocks/ops/binned_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu126-x86_64-linux/megablocks/ops/binned_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu126-x86_64-linux/megablocks/ops/gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu126-x86_64-linux/megablocks/ops/padded_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu126-x86_64-linux/megablocks/ops/padded_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu126-x86_64-linux/megablocks/ops/scatter.py

@@ -4,7 +4,7 @@
 from typing import Any, Optional
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx11-cu126-x86_64-linux/megablocks/ops/stk_autocast.py

@@ -0,0 +1,39 @@
+# vendored from
+# https://github.com/stanford-futuredata/stk/blob/736313768ef697ce13a0594a41b2512a0fbc9884/stk/backend/autocast.py
+import functools
+import torch
+
+
+def _is_eligible(x):
+    return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64)
+
+
+def _cast(x, dtype):
+    if isinstance(x, torch.Tensor) and _is_eligible(x):
+        return x.to(dtype)
+    elif isinstance(x, map):
+        return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()}
+    elif isinstance(x, list) or isinstance(x, tuple):
+        return type(x)(map(lambda y: _cast(y, dtype), x))
+    return x
+
+
+def custom_fwd(fwd):
+    """Wrap a custom autograd function that always uses autocast dtype."""
+
+    @functools.wraps(fwd)
+    def decorate_fwd(*args, **kwargs):
+        if torch.is_autocast_enabled():
+            with torch.autocast(device_type="cuda", enabled=False):
+                dtype = torch.get_autocast_gpu_dtype()
+                return fwd(*_cast(args, dtype), **_cast(kwargs, dtype))
+        return fwd(*args, **kwargs)
+    return decorate_fwd
+
+
+def custom_bwd(bwd):
+    @functools.wraps(bwd)
+    def decorate_bwd(*args, **kwargs):
+        with torch.autocast(device_type="cuda", enabled=False):
+            return bwd(*args, **kwargs)
+    return decorate_bwd

build/torch26-cxx98-cu118-x86_64-linux/megablocks/{_megablocks_dabb815.abi3.so → _megablocks_0586ba6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:defeb9a48abe98940478c79c5eac52f9fc7c22088abf9d119191559787bb95a9
+oid sha256:baacdb2bd8bcd004a86f63b0dc2754bac21214c9432bf6c00c464ccc26c25a83
 size 10510040

build/torch26-cxx98-cu118-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_dabb815
-ops = torch.ops._megablocks_dabb815
+from . import _megablocks_0586ba6
+ops = torch.ops._megablocks_0586ba6
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_dabb815::{op_name}"
+    return f"_megablocks_0586ba6::{op_name}"

build/torch26-cxx98-cu118-x86_64-linux/megablocks/ops/binned_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu118-x86_64-linux/megablocks/ops/binned_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu118-x86_64-linux/megablocks/ops/gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu118-x86_64-linux/megablocks/ops/padded_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu118-x86_64-linux/megablocks/ops/padded_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu118-x86_64-linux/megablocks/ops/scatter.py

@@ -4,7 +4,7 @@
 from typing import Any, Optional
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu118-x86_64-linux/megablocks/ops/stk_autocast.py

@@ -0,0 +1,39 @@
+# vendored from
+# https://github.com/stanford-futuredata/stk/blob/736313768ef697ce13a0594a41b2512a0fbc9884/stk/backend/autocast.py
+import functools
+import torch
+
+
+def _is_eligible(x):
+    return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64)
+
+
+def _cast(x, dtype):
+    if isinstance(x, torch.Tensor) and _is_eligible(x):
+        return x.to(dtype)
+    elif isinstance(x, map):
+        return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()}
+    elif isinstance(x, list) or isinstance(x, tuple):
+        return type(x)(map(lambda y: _cast(y, dtype), x))
+    return x
+
+
+def custom_fwd(fwd):
+    """Wrap a custom autograd function that always uses autocast dtype."""
+
+    @functools.wraps(fwd)
+    def decorate_fwd(*args, **kwargs):
+        if torch.is_autocast_enabled():
+            with torch.autocast(device_type="cuda", enabled=False):
+                dtype = torch.get_autocast_gpu_dtype()
+                return fwd(*_cast(args, dtype), **_cast(kwargs, dtype))
+        return fwd(*args, **kwargs)
+    return decorate_fwd
+
+
+def custom_bwd(bwd):
+    @functools.wraps(bwd)
+    def decorate_bwd(*args, **kwargs):
+        with torch.autocast(device_type="cuda", enabled=False):
+            return bwd(*args, **kwargs)
+    return decorate_bwd

build/torch26-cxx98-cu124-x86_64-linux/megablocks/{_megablocks_dabb815.abi3.so → _megablocks_0586ba6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b6cc8f982e35bfa07a9121807bacd3c3572d6ecb1495bcb2b6286b967fb20d58
+oid sha256:20a8e0a793ac29bc168d10e1c9e465082c2adb1582ff79d1a083798f9a955a5f
 size 11857920

build/torch26-cxx98-cu124-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_dabb815
-ops = torch.ops._megablocks_dabb815
+from . import _megablocks_0586ba6
+ops = torch.ops._megablocks_0586ba6
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_dabb815::{op_name}"
+    return f"_megablocks_0586ba6::{op_name}"

build/torch26-cxx98-cu124-x86_64-linux/megablocks/ops/binned_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu124-x86_64-linux/megablocks/ops/binned_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu124-x86_64-linux/megablocks/ops/gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu124-x86_64-linux/megablocks/ops/padded_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu124-x86_64-linux/megablocks/ops/padded_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu124-x86_64-linux/megablocks/ops/scatter.py

@@ -4,7 +4,7 @@
 from typing import Any, Optional
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu124-x86_64-linux/megablocks/ops/stk_autocast.py

@@ -0,0 +1,39 @@
+# vendored from
+# https://github.com/stanford-futuredata/stk/blob/736313768ef697ce13a0594a41b2512a0fbc9884/stk/backend/autocast.py
+import functools
+import torch
+
+
+def _is_eligible(x):
+    return x.is_floating_point() and x.is_cuda and (x.dtype is not torch.float64)
+
+
+def _cast(x, dtype):
+    if isinstance(x, torch.Tensor) and _is_eligible(x):
+        return x.to(dtype)
+    elif isinstance(x, map):
+        return {_cast(k, dtype): _cast(v, dtype) for k, v in x.items()}
+    elif isinstance(x, list) or isinstance(x, tuple):
+        return type(x)(map(lambda y: _cast(y, dtype), x))
+    return x
+
+
+def custom_fwd(fwd):
+    """Wrap a custom autograd function that always uses autocast dtype."""
+
+    @functools.wraps(fwd)
+    def decorate_fwd(*args, **kwargs):
+        if torch.is_autocast_enabled():
+            with torch.autocast(device_type="cuda", enabled=False):
+                dtype = torch.get_autocast_gpu_dtype()
+                return fwd(*_cast(args, dtype), **_cast(kwargs, dtype))
+        return fwd(*args, **kwargs)
+    return decorate_fwd
+
+
+def custom_bwd(bwd):
+    @functools.wraps(bwd)
+    def decorate_bwd(*args, **kwargs):
+        with torch.autocast(device_type="cuda", enabled=False):
+            return bwd(*args, **kwargs)
+    return decorate_bwd

build/torch26-cxx98-cu126-x86_64-linux/megablocks/{_megablocks_dabb815.abi3.so → _megablocks_0586ba6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:be9d5e5df42fd6d0db62397eae9c462b9775e952ce7f71fb687c3ea75dfe6a74
+oid sha256:b2f5209e69d36d632939923c20ab90c074fe0100d8a4efbabe5cdcd32ccbcfd2
 size 11923672

build/torch26-cxx98-cu126-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_dabb815
-ops = torch.ops._megablocks_dabb815
+from . import _megablocks_0586ba6
+ops = torch.ops._megablocks_0586ba6
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_dabb815::{op_name}"
+    return f"_megablocks_0586ba6::{op_name}"

build/torch26-cxx98-cu126-x86_64-linux/megablocks/ops/binned_gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu126-x86_64-linux/megablocks/ops/binned_scatter.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels
 

build/torch26-cxx98-cu126-x86_64-linux/megablocks/ops/gather.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from stk.backend.autocast import custom_bwd, custom_fwd
+from .stk_autocast import custom_bwd, custom_fwd
 
 from ..backend import kernels