C.W.K.
Stream
Lesson 06 of 08 · published

keras.ops: The Universal Layer

~10 min · backend

Level 0Keras Apprentice
0 XP0/97 lessons0/20 achievements
0/120 XP to next level120 XP to go0% complete

One namespace, every engine

keras.ops is the backbone of Keras 3's multi-backend architecture — and the part you actually touch when you write custom math. It implements the full NumPy API plus the neural-network primitives NumPy lacks (convolutions, activations, crossentropy), and dispatches each call to whatever backend is loaded. Deliberately mirroring NumPy's names and signatures is the design choice that makes it almost invisible: if you know NumPy, you already know most of this surface.

The four namespaces you'll reach for

NamespaceExamples
keras.ops (NumPy)matmul, sum, mean, reshape, concatenate, stack, einsum, arange, clip
keras.ops.nnsoftmax, sigmoid, relu, conv, depthwise_conv, binary_crossentropy
keras.ops.imageresize, crop_images, pad_images, rgb_to_grayscale
keras.randomnormal, uniform, dropout (with SeedGenerator)

One gotcha worth internalizing now: randomness lives in keras.random and is meant to be driven by a SeedGenerator, not the backend's global RNG. That's what keeps results reproducible and JAX-pure across all three engines.

Porting TF code, op for op

If you're migrating an existing TensorFlow custom layer, most of the work is a near-mechanical rename — same operation, NumPy-flavored name:

TensorFlowKeras 3
tf.reduce_sumkeras.ops.sum
tf.concatkeras.ops.concatenate
tf.rangekeras.ops.arange
tf.reduce_meankeras.ops.mean
tf.gatherkeras.ops.take
tf.clip_by_valuekeras.ops.clip

Do the rename once and the layer stops being TF-only — it now runs unchanged on PyTorch and JAX too, which is the principle the next lesson leans on.

External links

Exercise

Write a custom Layer subclass whose call() uses keras.ops.softmax + keras.ops.matmul. Run it under all three backends with no code change. Confirm outputs match within numerical tolerance.

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.