🔄 Serket + Keras

In this example, a simple serket model is converted and trained in the new keras with the jax backend.

!pip install git+https://github.com/ASEM000/serket --quiet
!pip install keras --quiet

Imports

import os

os.environ["KERAS_BACKEND"] = "jax"
import jax.numpy as jnp
import jax.tree_util as jtu
import jax.random as jr
import keras
import serket as sk
import jax

serket -> keras conversion

# define a keras layer that wraps the serket layer
def is_trainable(leaf):
    if isinstance(leaf, jax.Array) and jnp.issubdtype(leaf.dtype, jnp.inexact):
        return True
    return False


def serket_to_keras(layer: sk.TreeClass) -> keras.Layer:
    leaves, treedef = jax.tree_util.tree_flatten(sk.tree_mask(layer))

    class SerketToKeras(keras.Layer):
        def __init__(self, layer, name=None):
            """Converts a serket layer to a keras layer"""
            super().__init__(name=name)
            # extract the leaves from the serket layer
            # here leaves of a masked layer are the trainable variables
            # and treedef is the tree structure of the layer
            for leaf in leaves:
                variable = keras.Variable(
                    initializer=leaf, trainable=is_trainable(leaf)
                )
                self._track_variable(variable)
            # mark the layer as built
            self.built = True

        def call(self, x):
            """Applies the layer to the input"""
            # convert the keras variables to jax arrays to be used in serket
            leaves = jtu.tree_map(jnp.array, self.trainable_variables)
            # unflatten the layer with the updated leaves
            layer = jtu.tree_unflatten(treedef, leaves)
            # apply the layer after unmasking it
            return sk.tree_unmask(layer)(x)

        @property
        def model(self):
            leaves = jax.tree_map(jnp.array, self.trainable_variables)
            layer = jax.tree_util.tree_unflatten(treedef, leaves)
            return sk.tree_unmask(layer)

    return SerketToKeras(layer)

Define a simple serket layer

# lets define a simple model in serket
class Linear(sk.TreeClass):
    def __init__(self, in_features, out_features, *, key):
        k1, k2 = jr.split(key)
        self.in_features = in_features
        self.out_features = out_features
        self.weight = jr.normal(k1, (in_features, out_features))
        self.bias = jr.normal(k2, (out_features,))

    def __call__(self, x):
        return x @ self.weight + self.bias

Train in keras

sk_model = sk.Sequential(
    Linear(1, 20, key=jr.key(0)),
    jax.nn.tanh,
    Linear(20, 20, key=jr.key(1)),
    jax.nn.tanh,
)

# use serket with keras model
model = keras.Sequential([serket_to_keras(sk_model), keras.layers.Dense(1)])


model.compile(
    optimizer=keras.optimizers.Adam(1e-2),
    loss=keras.losses.MeanSquaredError(),
)

x = jnp.linspace(-1, 1, 100)[:, None]
y = x**2 + jr.normal(jr.key(0), (100, 1)) * 0.01
model.fit(x, y, epochs=100)

Epoch 1/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 72ms/step - loss: 1.3688
Epoch 2/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 482us/step - loss: 0.0929
Epoch 3/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 395us/step - loss: 0.2363
Epoch 4/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 389us/step - loss: 0.1136
Epoch 5/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 587us/step - loss: 0.0141
Epoch 6/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 378us/step - loss: 0.0579
Epoch 7/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 468us/step - loss: 0.0498
Epoch 8/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 380us/step - loss: 0.0191
Epoch 9/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 505us/step - loss: 0.0077
Epoch 10/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 356us/step - loss: 0.0138
Epoch 11/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 537us/step - loss: 0.0065
Epoch 12/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 438us/step - loss: 0.0031
Epoch 13/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 477us/step - loss: 0.0056
Epoch 14/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 382us/step - loss: 0.0042
Epoch 15/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 431us/step - loss: 0.0022
Epoch 16/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 344us/step - loss: 0.0021
Epoch 17/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 614us/step - loss: 0.0023
Epoch 18/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 598us/step - loss: 0.0015
Epoch 19/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 387us/step - loss: 0.0017
Epoch 20/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 675us/step - loss: 0.0015
Epoch 21/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 361us/step - loss: 0.0013
Epoch 22/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 400us/step - loss: 0.0011
Epoch 23/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 501us/step - loss: 0.0011
Epoch 24/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 507us/step - loss: 0.0010
Epoch 25/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 409us/step - loss: 9.4913e-04
Epoch 26/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 374us/step - loss: 8.5285e-04
Epoch 27/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 375us/step - loss: 9.5316e-04
Epoch 28/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 541us/step - loss: 7.9228e-04
Epoch 29/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 674us/step - loss: 7.8916e-04
Epoch 30/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 466us/step - loss: 8.7422e-04
Epoch 31/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 390us/step - loss: 8.2918e-04
Epoch 32/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 521us/step - loss: 8.0647e-04
Epoch 33/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 572us/step - loss: 8.8197e-04
Epoch 34/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 515us/step - loss: 6.1969e-04
Epoch 35/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 838us/step - loss: 7.9608e-04
Epoch 36/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 400us/step - loss: 6.4733e-04
Epoch 37/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 416us/step - loss: 5.9605e-04
Epoch 38/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 427us/step - loss: 6.3412e-04
Epoch 39/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 406us/step - loss: 6.9424e-04
Epoch 40/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 514us/step - loss: 4.8109e-04
Epoch 41/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 503us/step - loss: 7.7633e-04
Epoch 42/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 660us/step - loss: 6.0829e-04
Epoch 43/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 381us/step - loss: 0.0011
Epoch 44/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 359us/step - loss: 9.9451e-04
Epoch 45/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 417us/step - loss: 7.8189e-04
Epoch 46/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 942us/step - loss: 7.9073e-04
Epoch 47/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 944us/step - loss: 9.5854e-04
Epoch 48/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 464us/step - loss: 8.9555e-04
Epoch 49/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 569us/step - loss: 7.5290e-04
Epoch 50/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 661us/step - loss: 5.6784e-04
Epoch 51/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 1ms/step - loss: 6.1000e-04
Epoch 52/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 541us/step - loss: 4.6011e-04
Epoch 53/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 483us/step - loss: 3.7351e-04
Epoch 54/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 941us/step - loss: 3.9865e-04
Epoch 55/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 486us/step - loss: 4.7861e-04
Epoch 56/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 729us/step - loss: 3.9875e-04
Epoch 57/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 426us/step - loss: 2.7468e-04
Epoch 58/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 528us/step - loss: 4.4127e-04
Epoch 59/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 930us/step - loss: 3.3213e-04
Epoch 60/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 418us/step - loss: 3.3361e-04
Epoch 61/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 537us/step - loss: 3.3838e-04
Epoch 62/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 442us/step - loss: 3.0299e-04
Epoch 63/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 465us/step - loss: 2.9275e-04
Epoch 64/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 508us/step - loss: 2.8322e-04
Epoch 65/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 415us/step - loss: 2.8615e-04
Epoch 66/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 498us/step - loss: 2.5094e-04
Epoch 67/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 420us/step - loss: 2.6434e-04
Epoch 68/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 544us/step - loss: 3.1166e-04
Epoch 69/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 752us/step - loss: 2.4866e-04
Epoch 70/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 459us/step - loss: 2.3734e-04
Epoch 71/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 418us/step - loss: 2.0421e-04
Epoch 72/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 618us/step - loss: 2.4537e-04
Epoch 73/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 863us/step - loss: 2.4165e-04
Epoch 74/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 415us/step - loss: 2.3931e-04
Epoch 75/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 502us/step - loss: 2.3416e-04
Epoch 76/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 445us/step - loss: 2.5766e-04
Epoch 77/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 413us/step - loss: 2.0533e-04
Epoch 78/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 524us/step - loss: 2.7161e-04
Epoch 79/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 557us/step - loss: 2.8574e-04
Epoch 80/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 515us/step - loss: 1.8881e-04
Epoch 81/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 420us/step - loss: 2.0679e-04
Epoch 82/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 448us/step - loss: 2.0246e-04
Epoch 83/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 977us/step - loss: 2.0976e-04
Epoch 84/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 768us/step - loss: 2.0156e-04
Epoch 85/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 608us/step - loss: 1.8208e-04
Epoch 86/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 542us/step - loss: 1.8338e-04
Epoch 87/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 547us/step - loss: 1.7531e-04
Epoch 88/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 583us/step - loss: 1.7490e-04
Epoch 89/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 622us/step - loss: 1.5132e-04
Epoch 90/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 445us/step - loss: 1.6035e-04
Epoch 91/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 494us/step - loss: 1.5763e-04
Epoch 92/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 428us/step - loss: 1.4480e-04
Epoch 93/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 676us/step - loss: 1.6287e-04
Epoch 94/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 766us/step - loss: 1.5020e-04
Epoch 95/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 446us/step - loss: 1.5299e-04
Epoch 96/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 438us/step - loss: 1.7338e-04
Epoch 97/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 641us/step - loss: 1.5630e-04
Epoch 98/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 412us/step - loss: 1.3388e-04
Epoch 99/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 604us/step - loss: 1.4354e-04
Epoch 100/100
4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 797us/step - loss: 1.8138e-04

<keras.src.callbacks.history.History at 0x15db37ad0>

Extract trained layer

# take the serket layer from the keras layer at the end
model.layers[0].model

Sequential(
  layers=(
    Linear(
      in_features=1,
      out_features=20,
      weight=f32[1,20](μ=-0.19, σ=0.92, ∈[-1.65,1.91]),
      bias=f32[20](μ=-0.03, σ=0.95, ∈[-2.12,1.93])
    ),
    jit(tanh(x)),
    Linear(
      in_features=20,
      out_features=20,
      weight=f32[20,20](μ=-0.04, σ=0.96, ∈[-2.73,2.65]),
      bias=f32[20](μ=0.37, σ=0.82, ∈[-1.04,1.99])
    ),
    jit(tanh(x))
  )
)