Keras: one operator layer, explicit backend
Install omnibias-keras==0.0.2a1 and the backend you intend to use. Select the backend before importing Keras:
KERAS_BACKEND=torch python keras_layer.py
# Or KERAS_BACKEND=jax / KERAS_BACKEND=tensorflow
The layer surface differs from the network-jet APIs in Torch and JAX. This example checks output shape, finite values, and reconstruction from configuration.
keras_layer.py
"""Set KERAS_BACKEND=torch, jax, or tensorflow before starting Python."""
import keras
import numpy as np
from omnibias.keras import OMBU
layer = OMBU(num_channels=4, K=2, base="tanh")
x = keras.ops.convert_to_tensor(np.linspace(-1, 1, 32).reshape(8, 4).astype("float32"))
y = layer(x)
assert tuple(y.shape) == (8, 4)
assert np.isfinite(keras.ops.convert_to_numpy(y)).all()
config = layer.get_config()
reconstructed = OMBU.from_config(config)
assert tuple(reconstructed(x).shape) == (8, 4)
print("Backend:", keras.backend.backend(), "output:", tuple(y.shape))
See the Keras package for available operator layers and their configuration contracts.