keras

Provides a single high-level API for defining models and layers that transparently dispatches numerical computation to JAX, TensorFlow, PyTorch, or OpenVINO backends selected at import time via KERAS_BACKEND environment variable or ~/.keras/keras.json. The framework maintains a backend-agnostic source of truth in keras/src/ with generated public API surface in keras/api/, enabling seamless backend switching without code changes. Runtime dispatch follows two paths: symbolic execution during model construction (shape/dtype inference via compute_output_spec on KerasTensor objects) and eager execution during training/inference (forwarded to active backend implementation).

Defines neural network layers (Dense, Conv2D, LSTM, etc.) and operations (numpy-compatible ops, neural network ops, core backend ops) in keras/src/ that are completely decoupled from backend implementation. Each layer inherits from a base Layer class that implements compute_output_spec() for symbolic shape/dtype inference and call() for eager execution. Backend-specific implementations are injected at runtime through the active backend module, allowing the same layer code to execute on JAX, TensorFlow, PyTorch, or OpenVINO without modification.

Provides a callback system (keras/src/callbacks/) that enables monitoring and controlling training through hooks at various training stages: on_epoch_begin, on_epoch_end, on_batch_begin, on_batch_end, on_train_begin, on_train_end. Built-in callbacks include EarlyStopping (stop training when validation metric plateaus), ModelCheckpoint (save best model), ReduceLROnPlateau (reduce learning rate), TensorBoard (visualization), and CSVLogger (log metrics). Callbacks are executed synchronously during training and have access to training state (epoch, batch, metrics, model weights).

Provides a metric system (keras/src/metrics/) for computing and tracking statistics during training and evaluation. Metrics are stateful objects that accumulate values across batches and compute aggregate statistics (accuracy, AUC, precision, recall, etc.). Metrics are compiled into models via model.compile(metrics=[...]) and automatically computed during training/evaluation. The framework provides built-in metrics for classification, regression, and ranking tasks. Metrics support both eager and graph execution modes and work identically across all backends.

Provides optimizer implementations (keras/src/optimizers/) including SGD, Adam, RMSprop, and others that update model weights based on gradients. Optimizers are backend-agnostic and delegate gradient updates to backend-specific implementations. Learning rate scheduling is supported through LearningRateSchedule objects that adjust learning rate during training based on epoch or batch number. Optimizers support momentum, weight decay, gradient clipping, and other advanced features. All optimizers work identically across backends.

Provides loss functions (keras/src/losses/) for training objectives including classification losses (categorical_crossentropy, sparse_categorical_crossentropy), regression losses (mean_squared_error, mean_absolute_error), and ranking losses. Loss functions are compiled into models via model.compile(loss=...) and automatically computed during training. The framework automatically computes gradients with respect to loss using the active backend's autodiff system (JAX's jax.grad, PyTorch's autograd, TensorFlow's GradientTape). Loss computation and gradient backpropagation are handled transparently without user code.

Provides APIs for inspecting model structure and accessing weights: model.summary() displays layer structure and parameter counts, model.get_weights() returns all weights as NumPy arrays, model.set_weights() updates weights, model.get_config() returns model configuration as JSON, model.get_layer() retrieves specific layers by name. These APIs work identically across all backends and enable model analysis, weight manipulation, and configuration serialization without backend-specific code.

Exposes a NumPy-compatible operation API (keras.ops.numpy.*) that mirrors NumPy's function signatures and behavior while dispatching to backend-specific implementations. Operations include array manipulation (reshape, concatenate, transpose), mathematical functions (sin, exp, matmul), and linear algebra (linalg.solve, linalg.eigh). The dispatch mechanism routes each operation call to the active backend's implementation in keras/src/backend/{backend}/numpy.py, ensuring numerical consistency across backends while leveraging backend-specific optimizations.

Provides a comprehensive set of neural network operations (keras.ops.nn.*) including activations (relu, sigmoid, softmax), normalization (batch_norm, layer_norm), convolution, pooling, and attention mechanisms. These operations are implemented in keras/src/ops/nn.py with backend-specific implementations in keras/src/backend/{backend}/nn.py. Each operation supports automatic differentiation through the active backend's autodiff system (JAX's jax.grad, PyTorch's autograd, TensorFlow's GradientTape), enabling gradient computation for training without explicit implementation.

Provides a high-level training API (model.fit(), model.train_on_batch()) that abstracts away backend-specific training mechanics. The training loop handles gradient computation, optimizer updates, metric tracking, and callback execution in a backend-agnostic manner. Distributed training is supported through backend-specific mechanisms: JAX uses jax.experimental.multihost_utils, PyTorch uses torch.distributed, TensorFlow uses tf.distribute.Strategy. The framework automatically detects available devices and distributes computation across them without requiring user code changes.

Provides model serialization to multiple deployment formats: SavedModel (TensorFlow format), ONNX (framework-agnostic), LiteRT (mobile), and OpenVINO (edge inference). Export is handled through keras/src/saving/ with format-specific implementations. SavedModel export preserves the full model graph and weights; ONNX export converts the model to ONNX intermediate representation for cross-framework compatibility; LiteRT export optimizes for mobile devices; OpenVINO export targets edge devices. Each format supports different deployment scenarios and optimization levels.

Provides quantization capabilities (keras/src/quantization/) for reducing model size and inference latency through reduced precision (int8, float16). Quantization is applied through quantization policies that specify precision for weights, activations, and computations. The framework supports post-training quantization and quantization-aware training (QAT). Quantization is implemented in a backend-agnostic manner, with backend-specific optimizations for each framework (JAX uses jax.numpy operations, PyTorch uses torch.quantization, TensorFlow uses tf.quantization).

Provides DType (data type) policies that specify precision for layer computations, weights, and outputs. Policies enable mixed-precision training where computations use float32 for numerical stability but weights are stored in float16 for memory efficiency. DType policies are defined in keras/src/layers/layer.py and applied through layer.dtype_policy. The framework automatically handles precision conversion during forward/backward passes, leveraging backend-specific mixed-precision support (JAX's automatic mixed precision, PyTorch's autocast, TensorFlow's mixed_float16 policy).

Provides two high-level model definition APIs: Sequential API for simple linear stacks of layers, and Functional API for complex architectures with multiple inputs/outputs and skip connections. Both APIs are defined in keras/src/models/ and compile to the same underlying Model class. Sequential API uses list-based layer stacking; Functional API uses symbolic tensor composition where layer calls return KerasTensor objects that represent computation graph structure. Both APIs support the same training, evaluation, and export capabilities.

Provides a Subclassing API where developers inherit from keras.layers.Layer or keras.Model to implement custom layers and models with arbitrary Python logic. Subclassed layers override build() to create weights and call() to define forward computation. The framework automatically tracks weights, handles gradient computation, and manages state. Subclassing enables dynamic control flow (if/while statements), custom gradient computation (via @tf.custom_gradient), and arbitrary Python logic that cannot be expressed through Sequential or Functional APIs.