Repository: unslothai/unsloth
File: unsloth/kernels/utils.py (Python)

Code before:
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import triton
import ctypes
MAX_FUSED_SIZE : int = 65536
next_power_of_2 = triton.next_power_of_2
import functools
from unsloth import DEVICE_TYPE

# torch.cuda.amp.custom_fwd is deprecated >= 2.4
import torch
torch_Tensor = torch.Tensor
from packaging.version import Version
if DEVICE_TYPE == "cuda":
    if Version(torch.__version__) < Version("2.4.0"):
        torch_amp_custom_fwd = torch.cuda.amp.custom_fwd
        torch_amp_custom_bwd = torch.cuda.amp.custom_bwd
    else:
        torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "cuda")
        torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "cuda")
    pass
elif DEVICE_TYPE == "xpu":
    if Version(torch.__version__) < Version("2.4.0"):
        raise RuntimeError("torch.xpu currently only supports torch.version >= 2.6.0")
    else:
        torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "xpu")
        torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "xpu")


# tl.math.tanh now is libdevice.tanh
from packaging.version import Version
import triton
import triton.language as tl
if Version(triton.__version__) >= Version("3.0.0"):
    if DEVICE_TYPE == "xpu":
        triton_tanh = tl.extra.intel.libdevice.tanh
    else:
        from triton.language.extra import libdevice
        triton_tanh = libdevice.tanh
    triton_cast = tl.cast
else:
    triton_tanh = tl.math.tanh
    # No casting in old Triton versions
    @triton.jit
    def triton_cast(x, dtype):
        return x.to(dtype)
    pass
pass

Review comment (suggestion):
Should be ?
```suggestion
    if Version(torch.__version__) < Version("2.6.0"):
```

Code after:
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import triton
import ctypes
MAX_FUSED_SIZE : int = 65536
next_power_of_2 = triton.next_power_of_2
import functools
from typing import Optional
from unsloth import DEVICE_TYPE

# torch.cuda.amp.custom_fwd is deprecated >= 2.4
import torch
torch_Tensor = torch.Tensor
from packaging.version import Version

if DEVICE_TYPE == "xpu" and Version(torch.__version__) < Version("2.6.0"):
    raise RuntimeError("Intel xpu currently supports unsloth with torch.version >= 2.6.0")

if Version(torch.__version__) < Version("2.4.0"):
    torch_amp_custom_fwd = torch.cuda.amp.custom_fwd
    torch_amp_custom_bwd = torch.cuda.amp.custom_bwd
else:
    torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "cuda")
    torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "cuda")
pass

if DEVICE_TYPE == "xpu":
    torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "xpu")
    torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "xpu")


# tl.math.tanh now is libdevice.tanh
from packaging.version import Version
import triton
import triton.language as tl
if Version(triton.__version__) >= Version("3.0.0"):
    if DEVICE_TYPE == "xpu":
        triton_tanh = tl.extra.intel.libdevice.tanh
    else:
        from triton.language.extra import libdevice
        triton_tanh = libdevice.tanh
    triton_cast = tl.cast
else:
    triton_tanh = tl.math.tanh
    # No casting in old Triton versions
    @triton.jit
    def triton_cast(x, dtype):
        return x.to(dtype)