TopolearnCustom layers, losses and metrics stay portable across JAX, TensorFlow and PyTorch only if they use keras.ops.
`keras.ops` is Keras 3's portable tensor API. It implements most of the NumPy API (ops.sum, ops.matmul, ops.reshape, ops.where, ops.einsum...) plus neural-network functions (ops.softmax, ops.relu, ops.conv, ops.one_hot, ops.logsumexp) and core control flow (ops.cond, ops.while_loop, ops.scan). Each call dispatches to the active backend, so code written with it runs on JAX, TensorFlow and PyTorch alike.
Practical rules:
- Inside custom layers, losses and metrics, use
keras.opsinstead oftf.*,jnp.*ortorch.*, unless you deliberately target one backend. ops.convert_to_tensor(x)makes a backend tensor;ops.convert_to_numpy(t)brings it back.- Dtypes are strings:
"float32","int32","bfloat16".
Randomness goes through `keras.random`: normal, uniform, dropout, shuffle. Pass a seed or, better, a keras.random.SeedGenerator, which holds state as a variable so random ops behave correctly inside compiled and stateless (JAX) functions.
Because results are backend-native, printing them shows backend types; convert to NumPy for portable output.
Example
import keras
from keras import ops
x = ops.arange(6, dtype="float32")
m = ops.reshape(x, (2, 3))
print(ops.convert_to_numpy(ops.sum(m, axis=0))) # [3. 5. 7.]
print(ops.convert_to_numpy(ops.matmul(m, ops.transpose(m)))) # [[ 5. 14.] [14. 50.]]
print(ops.convert_to_numpy(ops.softmax(ops.array([0.0, 0.0])))) # [0.5 0.5]
print(ops.convert_to_numpy(ops.where(m > 2, m, 0.0))) # [[0. 0. 0.] [3. 4. 5.]]
print(ops.convert_to_numpy(ops.one_hot(ops.array([2, 0]), 3))) # [[0. 0. 1.] [1. 0. 0.]]
seed = keras.random.SeedGenerator(42)
noise = keras.random.normal((2, 3), seed=seed) # reproducible on every backend
print(tuple(noise.shape)) # (2, 3)Remember
keras.ops= NumPy API + NN ops, on every backend- Use it in custom layers, losses and metrics for portability
- Use
keras.randomwith a SeedGenerator for randomness
Watch out for
- Calling
tf.*functions in a custom layer, which breaks as soon as the backend is JAX or PyTorch.