API Reference Axon v0.4.1



A high-level interface for creating neural network models.

Activation functions.

Module for rendering various visual representations of Axon models.

Parameter initializers.

Functional implementations of common neural network layer operations.

Abstraction for modeling a reduction of a dataset with an accumulated state for a number of epochs.

Accumulated state in an Axon.Loop.

Implementations of loss-scalers for use in mixed precision training.

Loss functions.

Metric functions.

Utilities for creating mixed precision policies.

Represents a missing value of an optional node.

Implementations of common gradient-based optimization algorithms.

Parameter Schedules.

Container for returning stateful outputs from Axon layers.

Parameter update methods.