mod config;
pub use config::{AdamWOptions, FusedAdamWError, StepCoefficients, StepControl};
pub mod reference;
#[cfg(feature = "fused-adamw-device")]
mod kernel;
#[cfg(feature = "fused-adamw-device")]
mod device;
#[cfg(feature = "fused-adamw-device")]
pub mod storage;
pub mod stats_plan;
#[cfg(feature = "fused-adamw-device")]
pub use device::{AdamWState, AdamWUpdate, adamw_step};
#[cfg(test)]
mod existing_optimizer_tests;
#[cfg(feature = "gradient-guard")]
pub mod gradient_norm;
#[cfg(feature = "gradient-guard-device")]
mod guard_kernel;
#[cfg(feature = "gradient-guard-device")]
mod guard_device;
#[cfg(feature = "gradient-guard-device")]
pub use guard_device::{
AdamWEntry, GuardedAdamWUpdate, GradientStatsReport, SkipReason,
gradient_stats_sync, guarded_adamw_step,
};