einx 0.0.1__tar.gz

einx-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Florian Fervers
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

einx-0.0.1/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.1
+Name: einx
+Version: 0.0.1
+Summary: Tensor Operations Expressed in Einstein-Inspired Notation
+Home-page: https://github.com/fferflo/einx
+Author: Florian Fervers
+Author-email: florian.fervers@gmail.com
+License: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3
+License-File: LICENSE

einx-0.0.1/einx/__init__.py

@@ -0,0 +1,7 @@
+from . import param
+from .lru_cache import lru_cache
+from . import tree_util
+from . import expr
+from .op import *
+from . import nn
+from . import backend

einx-0.0.1/einx/backend/__init__.py

@@ -0,0 +1,73 @@
+import sys, einx
+
+backends = []
+backend_factories = {}
+
+from ._numpy import numpy
+backends.append(numpy)
+
+from ._jax import make_jax_backend
+backend_factories["jax"] = make_jax_backend
+
+from ._torch import make_torch_backend
+backend_factories["torch"] = make_torch_backend
+
+from ._tensorflow import make_tensorflow_backend
+backend_factories["tensorflow"] = make_tensorflow_backend
+
+from .tracer import tracer
+backends.append(tracer)
+
+def update():
+    for backend_name in list(backend_factories.keys()):
+        if backend_name in sys.modules:
+            backends.append(backend_factories[backend_name]())
+            del backend_factories[backend_name]
+update()
+
+
+type_to_backend = {}
+
+def _get1(tensor):
+    tensor_backend = type_to_backend.get(type(tensor), None)
+    if tensor_backend is None:
+        update()
+
+        if tensor_backend is None:
+            for tensor_backend in backends:
+                if isinstance(tensor, tensor_backend.tensor) and not isinstance(tensor, numpy.tensor):
+                    break
+            else:
+                # Default backend is numpy
+                tensor_backend = numpy
+
+        type_to_backend[type(tensor)] = tensor_backend
+    return tensor_backend
+
+def get(arg):
+    if isinstance(arg, str):
+        name = arg
+        for backend in backends:
+            if backend.name == name:
+                return backend
+        update()
+        for backend in backends:
+            if backend.name == name:
+                return backend
+        raise ValueError(f"Backend {name} not found")
+    else:
+        tensors = arg
+        if len(tensors) == 1:
+            return _get1(tensors[0])
+        backend = None
+        for tensor in tensors:
+            if not tensor is None:
+                backend2 = _get1(tensor)
+                if backend2 != numpy:
+                    if not backend is None and backend != backend2:
+                        raise ValueError(f"Got tensors with conflicting backends: {backend.__name__} and {backend2.__name__}")
+                    backend = backend2
+        if backend is None:
+            return numpy
+        else:
+            return backend

einx-0.0.1/einx/backend/_jax.py

@@ -0,0 +1,82 @@
+from functools import partial
+from .base import base_backend
+
+def make_jax_backend():
+    import jax as jax_
+    import jax.numpy as jnp
+    class jax(base_backend):
+        @staticmethod
+        def to_tensor(tensor):
+            return jnp.asarray(tensor)
+
+        tensor = jnp.ndarray
+        name = "jax"
+
+        cast = lambda tensor, dtype: tensor.astype(dtype)
+        reshape = jnp.reshape
+        transpose = jnp.transpose
+        broadcast_to = jnp.broadcast_to
+        einsum = partial(jnp.einsum, optimize="optimal")
+        dot = jnp.dot
+        swapaxes = jnp.swapaxes
+
+        stack = jnp.stack
+        concatenate = jnp.concatenate
+
+        zeros = jnp.zeros
+        ones = jnp.ones
+
+        add = jnp.add
+        subtract = jnp.subtract
+        multiply = jnp.multiply
+        true_divide = jnp.true_divide
+        floor_divide = jnp.floor_divide
+        divide = jnp.divide
+        logical_and = jnp.logical_and
+        logical_or = jnp.logical_or
+        where = jnp.where
+        less = jnp.less
+        less_equal = jnp.less_equal
+        greater = jnp.greater
+        greater_equal = jnp.greater_equal
+        equal = jnp.equal
+        not_equal = jnp.not_equal
+        maximum = jnp.maximum
+        minimum = jnp.minimum
+
+        sum = jnp.sum
+        mean = jnp.mean
+        var = jnp.var
+        std = jnp.std
+        prod = jnp.prod
+        count_nonzero = jnp.count_nonzero
+        any = jnp.any
+        all = jnp.all
+        min = jnp.amin
+        max = jnp.amax
+
+        def get_at(tensor, coordinates):
+            return tensor[coordinates]
+        def set_at(tensor, coordinates, updates):
+            return tensor.at[coordinates].set(updates)
+        def add_at(tensor, coordinates, updates):
+            return tensor.at[coordinates].add(updates)
+        def subtract_at(tensor, coordinates, updates):
+            return tensor.at[coordinates].add(-updates)
+
+        flip = jnp.flip
+        roll = jnp.roll
+
+        sqrt = jnp.sqrt
+        rsqrt = jax_.lax.rsqrt
+        square = jnp.square
+
+        allclose = jnp.allclose
+
+        vmap = jax_.vmap
+
+        class random:
+            def bernoulli(rng, p, shape):
+                return jax_.random.bernoulli(rng, p, shape)
+
+    return jax

einx-0.0.1/einx/backend/_numpy.py

@@ -0,0 +1,97 @@
+import numpy as np
+from functools import partial
+from .base import base_backend
+
+class numpy(base_backend):
+    @staticmethod
+    def to_tensor(tensor):
+        return np.asarray(tensor)
+
+    tensor = np.ndarray
+    name = "numpy"
+
+    cast = lambda tensor, dtype: tensor.astype(dtype)
+    reshape = np.reshape
+    transpose = np.transpose
+    broadcast_to = np.broadcast_to
+    einsum = partial(np.einsum, optimize="optimal")
+    dot = np.dot
+    swapaxes = np.swapaxes
+
+    stack = np.stack
+    concatenate = np.concatenate
+
+    zeros = np.zeros
+    ones = np.ones
+
+    add = np.add
+    subtract = np.subtract
+    multiply = np.multiply
+    true_divide = np.true_divide
+    floor_divide = np.floor_divide
+    divide = np.divide
+    logical_and = np.logical_and
+    logical_or = np.logical_or
+    where = np.where
+    less = np.less
+    less_equal = np.less_equal
+    greater = np.greater
+    greater_equal = np.greater_equal
+    equal = np.equal
+    not_equal = np.not_equal
+    maximum = np.maximum
+    minimum = np.minimum
+
+    sum = np.sum
+    mean = np.mean
+    var = np.var
+    std = np.std
+    prod = np.prod
+    count_nonzero = np.count_nonzero
+    any = np.any
+    all = np.all
+    min = np.amin
+    max = np.amax
+
+    def get_at(tensor, coordinates):
+        return tensor[coordinates]
+    def set_at(tensor, coordinates, updates):
+        tensor[coordinates] = updates
+        return tensor
+    def add_at(tensor, coordinates, updates):
+        tensor[coordinates] += updates
+        return tensor
+    def subtract_at(tensor, coordinates, updates):
+        tensor[coordinates] -= updates
+        return tensor
+
+    flip = np.flip
+    roll = np.roll
+
+    sqrt = np.sqrt
+    rsqrt = lambda x: 1.0 / np.sqrt(x)
+    square = np.square
+
+    allclose = np.allclose
+
+    def vmap(op, in_axes, out_axes):
+        if not isinstance(in_axes, (tuple, list)) or not isinstance(out_axes, (tuple, list)):
+            raise ValueError("in_axes and out_axes must be tuples or lists of integers")
+        def inner(*args):
+            if len(args) != len(in_axes):
+                raise ValueError(f"Expected {len(in_axes)} arguments, got {len(args)}")
+            value = set(arg.shape[axis] for arg, axis in zip(args, in_axes) if not axis is None)
+            if len(value) != 1:
+                raise ValueError(f"Expected all arguments to have same size along vmap axis, got {value}")
+            value = value.pop()
+            xs_stacks = [[]] * len(out_axes)
+            for i in range(value):
+                xs = op(*[arg[(slice(None),) * axis + (i,)] if not axis is None else arg for arg, axis in zip(args, in_axes)])
+                if len(xs) != len(out_axes):
+                    raise ValueError(f"Expected {len(out_axes)} arguments from vmapped function, got {len(xs)}")
+                for xs_stack, x in zip(xs_stacks, xs):
+                    xs_stack.append(x)
+            xs = tuple(np.stack(xs_stack, axis=out_axis) for out_axis, xs_stack in zip(out_axes, xs_stacks))
+            return xs
+        inner.__name__ = f"vmap({op.__name__ if '__name__' in dir(op) else str(op)}, in_axes={in_axes}, out_axes={out_axes})"
+        return inner

einx-0.0.1/einx/backend/_tensorflow.py

@@ -0,0 +1,119 @@
+from functools import partial
+from .base import base_backend
+
+def make_tensorflow_backend():
+    import tensorflow as tf
+    import tensorflow.experimental.numpy as tnp
+    class tensorflow(base_backend):
+        @staticmethod
+        def to_tensor(tensor):
+            tensor = tf.convert_to_tensor(tensor)
+            if any(s is None for s in tensor.shape):
+                raise ValueError("Tensorflow tensors with dynamic shape are not supported")
+            return tensor
+
+        tensor = tf.Tensor
+        name = "tensorflow"
+
+        cast = tf.cast
+        reshape = tf.reshape
+        transpose = tf.transpose
+        broadcast_to = tf.broadcast_to
+        einsum = partial(tnp.einsum, optimize="optimal")
+        dot = tnp.dot
+        swapaxes = tnp.swapaxes
+
+        stack = tnp.stack
+        concatenate = tnp.concatenate
+
+        zeros = lambda shape, dtype="float32": tf.zeros(shape, dtype=dtype)
+        ones = lambda shape, dtype="float32": tf.ones(shape, dtype=dtype)
+
+        add = tnp.add
+        subtract = tnp.subtract
+        multiply = tnp.multiply
+        true_divide = tnp.true_divide
+        floor_divide = tnp.floor_divide
+        divide = tnp.divide
+        logical_and = tnp.logical_and
+        logical_or = tnp.logical_or
+        where = tnp.where
+        less = tnp.less
+        less_equal = tnp.less_equal
+        greater = tnp.greater
+        greater_equal = tnp.greater_equal
+        equal = tnp.equal
+        not_equal = tnp.not_equal
+        maximum = tnp.maximum
+        minimum = tnp.minimum
+
+        sum = tnp.sum
+        mean = tnp.mean
+        var = tnp.var
+        var = tnp.std
+        prod = tnp.prod
+        count_nonzero = tnp.count_nonzero
+        any = tnp.any
+        all = tnp.all
+        min = tnp.min
+        max = tnp.max
+
+        def get_at(tensor, coordinates):
+            return tensor[coordinates]
+        def set_at(tensor, coordinates, updates):
+            return tf.tensor_scatter_nd_update(tensor, tf.stack(coordinates, axis=-1), updates)
+        def add_at(tensor, coordinates, updates):
+            return tf.tensor_scatter_nd_add(tensor, tf.stack(coordinates, axis=-1), updates)
+        def subtract_at(tensor, coordinates, updates):
+            return tf.tensor_scatter_nd_sub(tensor, tf.stack(coordinates, axis=-1), updates)
+
+        def flip(x, axis):
+            if isinstance(axis, int):
+                axis = [axis]
+            return tf.reverse(x, axis)
+        def roll(x, axis, shift):
+            if isinstance(axis, int):
+                axis = [axis]
+            if isinstance(shift, int):
+                shift = [shift]
+            return tf.roll(x, tuple(shift), axis=tuple(axis))
+
+        sqrt = tf.math.sqrt
+        rsqrt = tf.math.rsqrt
+        square = tnp.square
+
+        allclose = tnp.allclose
+
+        def vmap(op, in_axes, out_axes):
+            def inner(*args):
+                # TODO: suboptimal (?) implementation of vmap in tensorflow that transposes the vmapped axis to the front and calls tf.vectorized_map
+                if len(args) != len(in_axes):
+                    raise ValueError(f"Expected {len(in_axes)} arguments, got {len(args)}")
+                value = set(arg.shape[axis] for arg, axis in zip(args, in_axes) if not axis is None)
+                if len(value) != 1:
+                    raise ValueError(f"Expected all arguments to have same size along vmap axis, got {value}")
+                value = value.pop()
+
+                # Move vmapped axes to front
+                xs = []
+                for arg, axis in zip(args, in_axes):
+                    if not axis is None:
+                        if axis != 0:
+                            perm = [axis] + [a for a in range(len(arg.shape)) if a != axis]
+                            arg = tf.transpose(arg, perm=perm)
+                    else:
+                        arg = arg[tf.newaxis]
+                    xs.append(arg)
+
+                xs = tf.vectorized_map(lambda xs: op(*xs), xs)
+                if len(xs) != len(out_axes):
+                    raise ValueError(f"Expected {len(out_axes)} arguments from vmapped function, got {len(xs)}")
+
+                # Move vmapped axis to out_axis
+                xs = [tf.transpose(x, perm=[(a + 1 if a < out_axis else (0 if a == out_axis else a)) for a in range(len(x.shape))]) for x, out_axis in zip(xs, out_axes)]
+
+                return tuple(xs)
+            inner.__name__ = f"vmap({op.__name__ if '__name__' in dir(op) else str(op)}, in_axes={in_axes}, out_axes={out_axes})"
+            return inner
+
+    return tensorflow

einx-0.0.1/einx/backend/_torch.py

@@ -0,0 +1,124 @@
+import einx
+from .base import base_backend
+
+def to_tuple(x):
+    if isinstance(x, tuple):
+        return x
+    elif isinstance(x, list):
+        return tuple(x)
+    elif isinstance(x, np.ndarray):
+        return tuple(x.tolist())
+    else:
+        raise ValueError(f"Cannot convert {type(x)} to tuple")
+
+def make_torch_backend():
+    import torch as torch_
+    import torch._dynamo as _dynamo
+    class torch(base_backend):
+        @staticmethod
+        def to_tensor(tensor):
+            if torch_.is_tensor(tensor):
+                return tensor
+            else:
+                return torch_.asarray(tensor)
+
+        tensor = torch_.Tensor
+        name = "torch"
+
+        cast = lambda tensor, dtype: tensor.type(vars(torch_)[dtype] if isinstance(dtype, str) else dtype)
+        reshape = lambda tensor, shape: torch_.reshape(tensor, to_tuple(shape))
+        transpose = torch_.permute
+        broadcast_to = lambda tensor, shape: torch_.broadcast_to(tensor, to_tuple(shape))
+        einsum = torch_.einsum
+        dot = torch_.matmul
+        swapaxes = torch_.swapaxes
+
+        stack = torch_.stack
+        concatenate = torch_.cat
+
+        zeros = lambda shape, dtype="float32": torch_.zeros(*shape, dtype=vars(torch_)[dtype] if isinstance(dtype, str) else dtype)
+        ones = lambda shape, dtype="float32": torch_.ones(*shape, dtype=vars(torch_)[dtype] if isinstance(dtype, str) else dtype)
+
+        add = torch_.add
+        subtract = torch_.subtract
+        multiply = torch_.multiply
+        true_divide = torch_.true_divide
+        floor_divide = torch_.floor_divide
+        divide = torch_.divide
+        logical_and = torch_.logical_and
+        logical_or = torch_.logical_or
+        where = torch_.where
+        less = torch_.less
+        less_equal = torch_.less_equal
+        greater = torch_.greater
+        greater_equal = torch_.greater_equal
+        equal = torch_.equal
+        not_equal = torch_.not_equal
+        def maximum(a, b):
+            return torch_.maximum(torch.to_tensor(a), torch.to_tensor(b)) # TODO: add support for python scalars everywhere
+        def minimum(a, b):
+            return torch_.minimum(torch.to_tensor(a), torch.to_tensor(b))
+
+        sum = torch_.sum
+        mean = torch_.mean
+        var = torch_.var
+        std = torch_.std
+        prod = torch_.prod
+        count_nonzero = torch_.count_nonzero
+        any = torch_.any
+        all = torch_.all
+        min = torch_.min
+        max = torch_.max
+
+        def get_at(tensor, coordinates):
+            return tensor[coordinates]
+        def set_at(tensor, coordinates, updates):
+            tensor[coordinates] = updates
+            return tensor
+        def add_at(tensor, coordinates, updates):
+            tensor[coordinates] += updates
+            return tensor
+        def subtract_at(tensor, coordinates, updates):
+            tensor[coordinates] -= updates
+            return tensor
+
+        def flip(tensor, axis):
+            if isinstance(axis, int):
+                axis = [axis]
+            return torch_.flip(tensor, axis)
+        def roll(tensor, shift, axis):
+            if isinstance(axis, int):
+                axis = [axis]
+            return torch_.roll(tensor, shift, axis)
+
+        sqrt = torch_.sqrt
+        rsqrt = torch_.rsqrt
+        square = torch_.square
+
+        allclose = torch_.allclose
+
+        def vmap(op, in_axes, out_axes):
+            return torch_.vmap(
+                op,
+                in_dims=tuple(in_axes) if isinstance(in_axes, list) else in_axes,
+                out_dims=tuple(out_axes) if isinstance(out_axes, list) else out_axes,
+            )
+
+        class random:
+            def bernoulli(rng, p, shape):
+                return torch_.bernoulli(torch_.full(shape, p), generator=rng) > 0.5
+
+    _dynamo.allow_in_graph(einx.dot)
+    _dynamo.allow_in_graph(einx.rearrange)
+    _dynamo.allow_in_graph(einx.elementwise)
+    _dynamo.allow_in_graph(einx.reduce)
+    _dynamo.allow_in_graph(einx.vmap)
+    _dynamo.allow_in_graph(einx.vmap_with_axis)
+    _dynamo.allow_in_graph(einx.nn.norm)
+    _dynamo.allow_in_graph(einx.nn.linear)
+    _dynamo.allow_in_graph(einx.nn.dropout)
+
+    for op_name in einx.elementwise._op_names + einx.reduce._op_names + einx.vmap_with_axis._op_names:
+        _dynamo.allow_in_graph(getattr(einx, op_name))
+
+    return torch

einx-0.0.1/einx/backend/base.py

@@ -0,0 +1,18 @@
+import einx
+import numpy as np
+
+class base_backend:
+    @classmethod
+    def apply(backend, op, args, kwargs, output_shapes):
+        if isinstance(op, str):
+            x = backend
+            for name in op.split("."):
+                x = getattr(x, name)
+            op = x
+        result = op(*args, **kwargs)
+        def assert_shape(tensor, out_shape):
+            in_shape = np.asarray(tensor.shape)
+            out_shape = np.asarray(out_shape)
+            assert in_shape.shape == out_shape.shape and np.all(in_shape == out_shape), f"Expected shape {out_shape}, got {in_shape}"
+        einx.tree_util.tree_map(assert_shape, result, output_shapes)
+        return result