diff --git a/.github/workflows/ruff.yml b/.github/workflows/ruff.yml new file mode 100644 index 000000000..6d3602cd5 --- /dev/null +++ b/.github/workflows/ruff.yml @@ -0,0 +1,14 @@ +name: Ruff +on: [push, pull_request] +jobs: + ruff: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: chartboost/ruff-action@v1 + with: + src: './python/' + - uses: chartboost/ruff-action@v1 + with: + src: './python/' + args: format --check diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index 15f9ade74..8b11ca8a7 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -1,5 +1,6 @@ import contextlib +import infinicore.nn as nn from infinicore.device import device from infinicore.dtype import ( bfloat16, @@ -33,6 +34,7 @@ from infinicore.tensor import ( Tensor, empty, + empty_like, from_blob, ones, strided_empty, @@ -40,12 +42,13 @@ zeros, ) -from infinicore import nn as nn - __all__ = [ + # Modules. + "nn", # Classes. "device", "dtype", + "Tensor", # Data Types. "bfloat16", "bool", @@ -75,6 +78,7 @@ "matmul", "rearrange", "empty", + "empty_like", "from_blob", "ones", "strided_empty", diff --git a/python/infinicore/nn/functional.py b/python/infinicore/nn/functional.py index a49c88cb1..ea969052c 100644 --- a/python/infinicore/nn/functional.py +++ b/python/infinicore/nn/functional.py @@ -1,18 +1,15 @@ import infinicore from infinicore.lib import _infinicore +from infinicore.tensor import Tensor __all__ = ["causal_softmax", "rms_norm", "silu", "swiglu"] -def causal_softmax( - input: infinicore.Tensor, - out=None -) -> infinicore.Tensor: - r"""Apply a causal softmax function. - """ +def causal_softmax(input: Tensor, out=None) -> Tensor: + r"""Apply a causal softmax function.""" if out is None: - return infinicore.Tensor(_infinicore.causal_softmax(input._underlying)) + return Tensor(_infinicore.causal_softmax(input._underlying)) _infinicore.causal_softmax_(out._underlying, input._underlying) @@ -20,49 +17,50 @@ def causal_softmax( def rms_norm( - input: infinicore.Tensor, - normalized_shape: list[int], - weight: infinicore.Tensor, - eps: float = 1e-5, - out=None -) -> infinicore.Tensor: - r"""Apply Root Mean Square Layer Normalization. - """ - - assert normalized_shape == weight.shape, "normalized_shape does not match weight.shape." + input: Tensor, + normalized_shape: list[int], + weight: Tensor, + eps: float = 1e-5, + *, + out=None, +) -> Tensor: + r"""Apply Root Mean Square Layer Normalization.""" + + assert normalized_shape == weight.shape, ( + "normalized_shape does not match weight.shape." + ) if out is None: - return infinicore.Tensor( - _infinicore.rms_norm(input._underlying, weight._underlying, eps) - ) + return Tensor(_infinicore.rms_norm(input._underlying, weight._underlying, eps)) _infinicore.rms_norm_(out._underlying, input._underlying, weight._underlying, eps) return out -def silu(input: infinicore.Tensor, inplace: bool = False, out=None) -> infinicore.Tensor: - r"""Apply the Sigmoid Linear Unit (SiLU) function, element-wise. - """ +def silu(input: Tensor, inplace: bool = False, *, out=None) -> Tensor: + r"""Apply the Sigmoid Linear Unit (SiLU) function, element-wise.""" + + if infinicore.use_ntops and input.device.type in ("cuda", "musa") and out is None: + return infinicore.ntops.torch.silu(input, inplace=inplace) if inplace: _infinicore.silu_(input._underlying, input._underlying) return input if out is None: - return infinicore.Tensor(_infinicore.silu(input._underlying)) + return Tensor(_infinicore.silu(input._underlying)) _infinicore.silu_(out._underlying, input._underlying) return out -def swiglu(input: infinicore.Tensor, other: infinicore.Tensor, out=None): - r"""Apply the Swish-Gated Linear Unit (SwiGLU) function, element-wise. - """ +def swiglu(input: Tensor, other: Tensor, *, out=None): + r"""Apply the Swish-Gated Linear Unit (SwiGLU) function, element-wise.""" if out is None: - return infinicore.Tensor(_infinicore.swiglu(input._underlying, other._underlying)) + return Tensor(_infinicore.swiglu(input._underlying, other._underlying)) _infinicore.swiglu_(out._underlying, input._underlying, other._underlying) diff --git a/python/infinicore/tensor.py b/python/infinicore/tensor.py index d40d8d10d..b4eada011 100644 --- a/python/infinicore/tensor.py +++ b/python/infinicore/tensor.py @@ -77,7 +77,7 @@ def view(self, shape): def debug(self, filename=None): """Print tensor data or save to file for debugging - + Args: filename: Optional filename to save raw binary data. If None, prints to stdout. """ @@ -93,6 +93,16 @@ def empty(size, *, dtype=None, device=None, pin_memory=False): ) +def empty_like(input, *, dtype=None, device=None): + if dtype is None: + dtype = input.dtype + + if device is None: + device = input.device + + return empty(input.size(), dtype=dtype, device=device) + + def strided_empty(size, strides, *, dtype=None, device=None, pin_memory=False): return Tensor( _infinicore.strided_empty(