da4ml 0.5.0__cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

da4ml/converter/hgq2/layers/conv.py

@@ -0,0 +1,149 @@
+from typing import TypeVar
+
+import numpy as np
+from hgq.layers import (
+    QConv1D,
+    QConv2D,
+    QConv3D,
+)
+from keras import ops
+from keras.src.ops.image import ExtractPatches, extract_patches_3d
+
+from ....trace import FixedVariableArray
+from ._base import ReplayOperationBase, to_np_arr
+
+T = TypeVar('T', FixedVariableArray, np.ndarray)
+
+
+def symbolic_extract_patches_3d(
+    images: T,
+    size: tuple[int, int, int],
+    strides: tuple[int, int, int],
+    dilation_rate: tuple[int, int, int],
+    padding: str,
+    data_format: str,
+    pad_value: float = 0,
+) -> T:
+    img_tensor = ops.cast(ops.reshape(ops.arange(images.size), images.shape), dtype='float32')
+    img_tensor = -img_tensor - 1
+    out_tensor = extract_patches_3d(
+        img_tensor[None],
+        size=size,
+        strides=strides,
+        dilation_rate=dilation_rate,  # type: ignore
+        padding=padding,
+        data_format=data_format,
+    )[0]
+    out_index: np.ndarray = ops.convert_to_numpy(out_tensor).round().astype(np.int32)  # type: ignore
+    mask = out_index == 0
+    out_index = np.where(mask, 0, -out_index - 1)
+    images = images.ravel()[out_index]
+
+    if isinstance(images, FixedVariableArray):
+        _vars = images._vars
+        _vars = np.where(mask, pad_value, _vars)
+        images = FixedVariableArray(_vars, images.solver_options)  # type: ignore
+    else:
+        images = np.where(mask, pad_value, images)
+
+    return images
+
+
+def symbolic_extract_patches(
+    images: T,
+    size: tuple[int, ...] | int,
+    strides: tuple[int, ...] | int,
+    dilation_rate: tuple[int, ...] | int,
+    padding: str,
+    data_format: str,
+    pad_value: float = 0,
+) -> T:
+    rank = images.ndim - 1
+    size = (size,) * rank if isinstance(size, int) else size
+    strides = (strides,) * rank if isinstance(strides, int) else strides
+    dilation_rate = (dilation_rate,) * rank if isinstance(dilation_rate, int) else dilation_rate
+
+    assert rank == len(size) == len(strides) == len(dilation_rate), (
+        f'Invalid rank {rank} for size {size}, strides {strides}, dilation_rate {dilation_rate}'
+    )
+
+    pad_rank = 3 - rank
+    _size: tuple[int, int, int] = (1,) * pad_rank + size  # type: ignore
+    _strides: tuple[int, int, int] = (1,) * pad_rank + strides  # type: ignore
+    _dilation_rate: tuple[int, int, int] = (1,) * pad_rank + dilation_rate  # type: ignore
+
+    _pad = (1,) * pad_rank
+    if data_format == 'channels_first':
+        images = np.moveaxis(images, 0, -1)  # type: ignore
+
+    *spa, ch = images.shape
+    images = images.reshape(*_pad, *spa, ch)
+
+    r = symbolic_extract_patches_3d(
+        images,
+        size=_size,
+        strides=_strides,
+        dilation_rate=_dilation_rate,
+        padding=padding,
+        data_format='channels_last',
+        pad_value=pad_value,
+    )
+
+    return r.reshape(r.shape[pad_rank:])
+
+
+class ReplayExtractPatches(ReplayOperationBase):
+    handles = (ExtractPatches,)
+
+    def call(self, images: FixedVariableArray) -> FixedVariableArray:
+        op: ExtractPatches = self.op
+        pixel_shape = op.size
+        strides = op.strides
+        dilation_rate: int | tuple[int, int] = op.dilation_rate
+        padding = op.padding
+        data_format = op.data_format
+
+        if strides is None:
+            strides = 1
+
+        return symbolic_extract_patches(images, pixel_shape, strides, dilation_rate, padding, data_format)
+
+
+class ReplayQConv(ReplayOperationBase):
+    handles = (QConv1D, QConv2D, QConv3D)
+
+    def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
+        layer: QConv1D | QConv2D | QConv3D = self.op
+        qkernel = to_np_arr(layer.qkernel)
+        qbias = to_np_arr(layer.qbias) if layer.qbias is not None else None
+        strides = layer.strides
+        padding = layer.padding
+        dilation_rate = layer.dilation_rate
+        groups = layer.groups
+
+        if layer.data_format == 'channels_first':
+            inputs = np.moveaxis(inputs, 0, -1)  # type: ignore
+
+        x = symbolic_extract_patches(
+            inputs,
+            size=layer.kernel_size,
+            strides=strides,
+            dilation_rate=dilation_rate,
+            padding=padding,
+            data_format=layer.data_format,
+        )
+
+        ch_out = qkernel.shape[-1]
+
+        _ch_out = ch_out // groups
+
+        x = x.reshape(*x.shape[:-1], -1, groups)
+        kernel = qkernel.reshape(-1, groups, _ch_out)
+
+        outputs = np.einsum('...ig,igo->...go', x, kernel)  # type: ignore
+        outputs = outputs.reshape(*outputs.shape[:-2], -1) + qbias
+
+        if layer.data_format == 'channels_first':
+            outputs: FixedVariableArray = np.moveaxis(outputs, -1, 0)  # type: ignore
+
+        return outputs

da4ml/converter/hgq2/layers/dense.py

@@ -0,0 +1,39 @@
+import keras
+from hgq.layers import (
+    QBatchNormDense,
+    QDense,
+    QEinsumDense,
+    QEinsumDenseBatchnorm,
+)
+
+from ....trace import FixedVariableArray
+from ....trace.ops import einsum
+from ._base import ReplayOperationBase, to_np_arr
+
+
+class ReplayQDense(ReplayOperationBase):
+    handles = (QDense, QEinsumDense, QEinsumDenseBatchnorm, QBatchNormDense, keras.layers.EinsumDense)
+
+    def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
+        op = self.op
+        if isinstance(op, (QDense, QBatchNormDense)):
+            qkernel = op.qkernel
+            qbias = op.qbias
+            eq = '...c,cC->...C'
+        elif isinstance(op, (QEinsumDense, QEinsumDenseBatchnorm)):
+            qkernel = op.qkernel
+            qbias = op.qbias
+            eq = op.equation
+        elif isinstance(op, keras.layers.EinsumDense):
+            qkernel = op.kernel
+            qbias = op.bias
+            eq = op.equation
+        else:
+            raise TypeError(f'Unsupported layer type: {type(op)}')
+
+        qkernel = to_np_arr(qkernel)
+        qbias = to_np_arr(qbias) if qbias is not None else None
+        return (einsum(eq, inputs[None], qkernel) + qbias)[0]
+
+
+__all__ = ['ReplayQDense']

da4ml/converter/hgq2/layers/ops.py

@@ -0,0 +1,240 @@
+from collections.abc import Sequence
+
+import keras
+import numpy as np
+from hgq.layers import (
+    QAdd,
+    QAveragePow2,
+    QDot,
+    QEinsum,
+    QMaximum,
+    QMeanPow2,
+    QMinimum,
+    QMultiply,
+    QSubtract,
+    QSum,
+)
+from keras.src.ops.numpy import (
+    Abs,
+    Absolute,
+    Add,
+    Concatenate,
+    Divide,
+    Dot,
+    Einsum,
+    GetItem,
+    Matmul,
+    Max,
+    Maximum,
+    Min,
+    Minimum,
+    Moveaxis,
+    Multiply,
+    Ravel,
+    Repeat,
+    Reshape,
+    Subtract,
+    Sum,
+    Transpose,
+    TrueDivide,
+)
+
+from ....trace import FixedVariableArray
+from ....trace.ops import einsum
+from ._base import ReplayOperationBase
+
+
+class ReplayReshape(ReplayOperationBase):
+    handles = (keras.layers.Reshape, keras.layers.Flatten, Reshape, Ravel)
+
+    def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
+        if isinstance(self.op, (keras.layers.Flatten, Ravel)):
+            return inputs.ravel()
+        elif isinstance(self.op, keras.layers.Reshape):
+            return inputs.reshape(self.op.target_shape)
+        elif isinstance(self.op, Reshape):
+            return inputs.reshape(self.op.newshape[1:])
+        else:
+            raise TypeError(f'Unsupported layer type: {type(self.op)}')
+
+
+class ReplayMerge(ReplayOperationBase):
+    handles = (keras.layers.Add, keras.layers.Concatenate, QAdd, QMultiply, QSubtract, QMaximum, QMinimum, QAveragePow2)
+
+    def call(self, inputs: tuple[FixedVariableArray, ...]) -> FixedVariableArray:
+        op = self.op
+        name = op.__class__.__name__
+        if name.startswith('Q'):
+            name = name[1:]
+        _inputs: FixedVariableArray = np.stack(np.broadcast_arrays(*inputs), axis=0)  # type: ignore
+        match name:
+            case 'Add':
+                return np.sum(_inputs, axis=0)  # type: ignore
+            case 'AveragePow2':
+                return np.sum(_inputs, axis=0) * op._scale  # type: ignore
+            case 'Subtract':
+                assert len(_inputs) == 2, 'Subtract operation requires exactly two inputs'
+                return _inputs[0] - _inputs[1]
+            case 'Multiply':
+                return np.prod(_inputs, axis=0)  # type: ignore
+            case 'Maximum':
+                return np.amax(_inputs, axis=0)  # type: ignore
+            case 'Minimum':
+                return np.amin(_inputs, axis=0)  # type: ignore
+            case 'Concatenate':
+                return np.concatenate(_inputs, axis=op.axis)  # type: ignore
+
+            case _:
+                raise TypeError(f'Unsupported layer type: {type(op)}')
+
+
+class ReplayRepeatVector(ReplayOperationBase):
+    handles = (keras.layers.RepeatVector,)
+
+    def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
+        layer: keras.layers.RepeatVector = self.op
+        if layer.n == 1:
+            return inputs
+        # return FixedVariableArray(np.repeat(inputs._vars, layer.n, axis=0), inputs.solver_options)
+        return np.repeat(inputs[None], layer.n, axis=0)[0]  # type: ignore
+
+
+class ReplayGetItem(ReplayOperationBase):
+    handles = (GetItem,)
+
+    def call(self, x: FixedVariableArray, key):
+        if isinstance(key, list):
+            key = tuple(key)
+        return x[None][key][0]
+
+
+class ReplayReduction(ReplayOperationBase):
+    handles = (Sum, Max, Min)
+
+    def call(self, x: FixedVariableArray, axis=None, keepdims=False):
+        if isinstance(self.op, Sum):
+            op = np.sum
+        elif isinstance(self.op, Max):
+            op = np.amax
+        elif isinstance(self.op, Min):
+            op = np.amin
+        return op(x[None], axis=axis, keepdims=keepdims)[0]  # type: ignore
+
+
+class ReplayQReduction(ReplayOperationBase):
+    handles = (QSum, QMeanPow2)
+
+    def call(self, x: FixedVariableArray):
+        layer: QSum = self.op
+        axes, scale, keepdims = layer.axes, layer.scale, layer.keepdims
+        return np.sum(x[None], axis=axes, keepdims=keepdims)[0] * scale  # type: ignore
+
+
+class ReplayArithmetic(ReplayOperationBase):
+    handles = (Add, Subtract, Multiply, TrueDivide, Divide, Maximum, Minimum)
+
+    def call(self, x1: FixedVariableArray, x2: FixedVariableArray):
+        name = self.op.__class__.__name__
+        match name:
+            case 'Add':
+                return x1 + x2
+            case 'Subtract':
+                return x1 - x2
+            case 'Multiply':
+                return x1 * x2
+            case 'TrueDivide' | 'Divide':
+                return x1 / x2
+            case 'Maximum':
+                return np.maximum(x1, x2)  # type: ignore
+            case 'Minimum':
+                return np.minimum(x1, x2)  # type: ignore
+            case _:
+                raise TypeError(f'Unsupported arithmetic operation: {type(self.op)}')
+
+
+class ReplayConcatenate(ReplayOperationBase):
+    handles = (Concatenate,)
+
+    def call(self, xs: Sequence[FixedVariableArray]):
+        axis = self.op.axis
+        # return backend.numpy.concatenate(xs, axis=self.axis)
+        # return FixedVariableArray(np.concatenate([x._vars[None] for x in xs], axis=axis)[0], xs[0].solver_options)
+        return np.concatenate([x[None] for x in xs], axis=axis)[0]  # type: ignore
+
+
+class ReplayRepeat(ReplayOperationBase):
+    handles = (Repeat,)
+
+    def call(self, x: FixedVariableArray):
+        repeats, axis = self.op.repeats, self.op.axis
+        # return FixedVariableArray(np.repeat(x._vars[None], repeats, axis=axis)[0], x.solver_options)
+        return np.repeat(x[None], repeats, axis=axis)[0]  # type: ignore
+
+
+class ReplayTranspose(ReplayOperationBase):
+    handles = (Transpose,)
+
+    def call(self, x: FixedVariableArray):
+        axes = self.op.axes
+        return np.transpose(x, axes)  # type: ignore
+
+
+class ReplayMoveaxis(ReplayOperationBase):
+    handles = (Moveaxis,)
+
+    def call(self, x: FixedVariableArray):
+        source, destination = self.op.source, self.op.destination
+        return np.moveaxis(x[None], source, destination)[0]  # type: ignore
+
+
+class ReplayNoOp(ReplayOperationBase):
+    __noop_layers = []
+    for k, v in keras.layers.__dict__.items():
+        name = k.lower()
+        if 'dropout' in name or 'random' in name or 'noise' in name:
+            __noop_layers.append(v)
+
+    handles = tuple(__noop_layers)
+
+    def call(self, x: FixedVariableArray, training=False) -> FixedVariableArray:
+        assert not training, 'Training mode is not supported in mirror operation'
+        return x
+
+
+class ReplayEinsum(ReplayOperationBase):
+    handles = (QEinsum, Einsum, QDot, keras.layers.Dot)
+
+    def call(self, inputs: tuple[FixedVariableArray, FixedVariableArray]) -> FixedVariableArray:
+        op = self.op
+        if isinstance(op, QEinsum):
+            eq = op.equation
+        elif isinstance(op, Einsum):
+            eq = op.subscripts
+        else:  # QDot/Dot
+            dim0, dim1 = inputs[0].ndim + 1, inputs[1].ndim + 1
+            letters = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'[0 : dim0 + dim1]
+            sub0, sub1 = letters[:dim0], letters[dim0 : dim0 + dim1]
+            axes = list(op.axes) if not isinstance(op.axes, int) else [op.axes, op.axes]
+            idx0, idx1 = axes[0] if axes[0] >= 0 else axes[0] % dim0, axes[1] if axes[1] >= 0 else axes[1] % dim1
+            sub1 = sub1[:idx1] + sub0[idx0] + sub1[idx1 + 1 :]
+            sub_out = list(sub0 + sub1)
+            sub_out.remove(sub0[idx0])
+            sub_out.remove(sub0[idx0])
+            sub_out = ''.join(sub_out)
+            eq = f'{sub0},{sub1}->{sub_out}'
+        assert len(inputs) == 2, 'Only (Q)Einsum operations with exactly two inputs are supported'
+        return einsum(eq, inputs[0][None], inputs[1][None])[0]
+
+
+class ReplayMatmul(ReplayOperationBase):
+    handles = (Matmul, Dot)
+
+    def call(self, x1: FixedVariableArray, x2: FixedVariableArray) -> FixedVariableArray:
+        return x1 @ x2
+
+
+class ReplayAbs(ReplayOperationBase):
+    handles = (Absolute, Abs)
+
+    def call(self, x: FixedVariableArray) -> FixedVariableArray:
+        return np.abs(x)  # type: ignore

da4ml/converter/hgq2/layers/pool.py

@@ -0,0 +1,107 @@
+from math import prod
+
+import hgq
+import keras
+import numpy as np
+from keras.src.layers.pooling.base_global_pooling import BaseGlobalPooling
+from keras.src.layers.pooling.base_pooling import BasePooling
+
+from ....trace import FixedVariableArray
+from ._base import ReplayOperationBase
+from .conv import symbolic_extract_patches
+
+
+class ReplayPool(ReplayOperationBase):
+    handles = (
+        hgq.layers.QAvgPool1D,
+        hgq.layers.QAvgPool2D,
+        hgq.layers.QAvgPool3D,
+        hgq.layers.QMaxPool1D,
+        hgq.layers.QMaxPool2D,
+        hgq.layers.QMaxPool3D,
+        hgq.layers.QGlobalAveragePooling1D,
+        hgq.layers.QGlobalMaxPooling1D,
+        hgq.layers.QGlobalAveragePooling2D,
+        hgq.layers.QGlobalMaxPooling2D,
+        hgq.layers.QGlobalAveragePooling3D,
+        hgq.layers.QGlobalMaxPooling3D,
+        keras.layers.AveragePooling1D,
+        keras.layers.AveragePooling2D,
+        keras.layers.AveragePooling3D,
+        keras.layers.MaxPooling1D,
+        keras.layers.MaxPooling2D,
+        keras.layers.MaxPooling3D,
+        keras.layers.GlobalAveragePooling1D,
+        keras.layers.GlobalMaxPooling1D,
+        keras.layers.GlobalAveragePooling2D,
+        keras.layers.GlobalMaxPooling2D,
+        keras.layers.GlobalAveragePooling3D,
+        keras.layers.GlobalMaxPooling3D,
+    )
+
+    def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
+        cname = self.op.__class__.__name__
+        if 'Max' in cname:
+            op = 'max'
+        else:
+            assert 'Average' in cname, f'Unsupported global pooling layer: {cname}'
+            op = 'avg'
+
+        data_format = self.op.data_format
+        if data_format == 'channels_first':
+            inputs = np.moveaxis(inputs, 1, -1)  # type: ignore
+
+        if isinstance(self.op, BaseGlobalPooling):
+            pool_dim = self.op.input_spec.ndim - 2  # type: ignore
+            axis = tuple(range(pool_dim))
+            keepdims = self.op.keepdims
+
+            if op == 'max':
+                out = np.amax(inputs, axis=axis, keepdims=keepdims)  # type: ignore
+            elif op == 'avg':
+                pool_size = prod(inputs.shape[:-1])
+                out = np.sum(inputs, axis=axis, keepdims=keepdims) / pool_size  # type: ignore
+        else:
+            assert isinstance(self.op, BasePooling), f'Unknown pooling layer: {type(self.op)}'
+            pool_size = self.op.pool_size
+            strides = self.op.strides
+            padding = self.op.padding
+            pool_dim = len(pool_size)
+            ch = inputs.shape[-1]
+            x = symbolic_extract_patches(
+                inputs,
+                pool_size,
+                strides,
+                dilation_rate=1,
+                padding=padding,
+                data_format='channels_last',
+            )
+            x = x.reshape(x.shape[:-1] + (-1, ch))
+
+            if padding == 'same':
+                mask = symbolic_extract_patches(
+                    np.ones(inputs.shape, dtype=np.int32),
+                    pool_size,
+                    strides,
+                    dilation_rate=1,
+                    padding=padding,
+                    data_format='channels_last',
+                ).reshape(x.shape)
+            elif padding == 'valid':
+                mask = np.ones(x.shape, dtype=np.int32)
+            else:
+                raise ValueError(f'Unknown padding type: {padding}')
+
+            if op == 'max':
+                _vars = np.where(mask, x._vars, -(65535**2))
+                x = FixedVariableArray(_vars, x.solver_options)
+                out = np.max(x, axis=-2)  # type: ignore
+            elif op == 'avg':
+                out = np.sum(x, axis=-2) / np.sum(mask, axis=-2)  # type: ignore
+            else:
+                raise ValueError(f'Unknown pooling operation: {op}')
+
+        if data_format == 'channels_first':
+            out = np.moveaxis(out, -1, 1)  # type: ignore
+
+        return out  # type: ignore