1#![allow(clippy::type_complexity)]
2
3pub mod layers;
38pub mod loss;
39pub mod models;
40pub mod optimizers;
41pub mod persistence;
42pub mod random;
43pub mod schedulers;
44pub mod text;
45pub mod training;
46pub mod utils;
48
49pub use layers::bilstm_network::{BiLSTMNetwork, BiLSTMNetworkCache, CombineMode};
51pub use layers::dropout::{Dropout, Zoneout};
52pub use layers::gru_cell::{GRUCell, GRUCellCache, GRUCellGradients};
53pub use layers::linear::{LinearGradients, LinearLayer};
54pub use layers::lstm_cell::{LSTMCell, LSTMCellCache, LSTMCellGradients};
55pub use layers::peephole_lstm_cell::{
56 PeepholeLSTMCell, PeepholeLSTMCellCache, PeepholeLSTMCellGradients,
57};
58pub use loss::{CrossEntropyLoss, LossFunction, MAELoss, MSELoss};
59pub use models::gru_network::{
60 GRUNetwork, GRUNetworkCache, LayerDropoutConfig as GRULayerDropoutConfig,
61};
62pub use models::lstm_network::{LSTMNetwork, LSTMNetworkCache, LSTMState, LayerDropoutConfig};
63pub use optimizers::{Adam, RMSprop, ScheduledOptimizer, SGD};
64pub use persistence::{ModelMetadata, ModelPersistence, PersistenceError, PersistentModel};
65pub use random::seed;
66pub use schedulers::{
67 AnnealStrategy, ConstantLR, CosineAnnealingLR, CosineAnnealingWarmRestarts, CyclicalLR,
68 CyclicalMode, ExponentialLR, LRScheduleVisualizer, LearningRateScheduler, LinearLR,
69 MultiStepLR, OneCycleLR, PolynomialLR, ReduceLROnPlateau, ScaleMode, StepLR, WarmupScheduler,
70};
71pub use text::{
72 argmax, sample_nucleus, sample_top_k, sample_with_temperature, softmax, CharacterEmbedding,
73 EmbeddingGradients, TextVocabulary,
74};
75pub use training::{
76 create_adam_batch_trainer, create_basic_batch_trainer, create_basic_trainer,
77 create_cosine_annealing_trainer, create_one_cycle_trainer, create_step_lr_trainer,
78 EarlyStopper, EarlyStoppingConfig, EarlyStoppingMetric, LSTMBatchTrainer, LSTMTrainer,
79 ScheduledLSTMTrainer, TrainingConfig, TrainingMetrics,
80};
81
82#[cfg(test)]
83mod tests {
84 use super::*;
85 use ndarray::arr2;
86
87 #[test]
88 fn test_library_integration() {
89 let mut network = models::lstm_network::LSTMNetwork::new(2, 3, 1);
90 let input = arr2(&[[1.0], [0.5]]);
91 let state = network.zero_state(1);
92
93 let (hy, next) = network.forward(&input, &state);
94 let cy = &next.c[0];
95
96 assert_eq!(hy.shape(), &[3, 1]);
97 assert_eq!(cy.shape(), &[3, 1]);
98 }
99
100 #[test]
101 fn test_library_with_dropout() {
102 let mut network = models::lstm_network::LSTMNetwork::new(2, 3, 1)
103 .with_input_dropout(0.2, false)
104 .with_recurrent_dropout(0.3, true)
105 .with_output_dropout(0.1);
106
107 let input = arr2(&[[1.0], [0.5]]);
108 let state = network.zero_state(1);
109
110 network.train();
112 let (hy_train, next_train) = network.forward(&input, &state);
113 let cy_train = &next_train.c[0];
114
115 network.eval();
117 let (hy_eval, next_eval) = network.forward(&input, &state);
118 let cy_eval = &next_eval.c[0];
119
120 assert_eq!(hy_train.shape(), &[3, 1]);
121 assert_eq!(cy_train.shape(), &[3, 1]);
122 assert_eq!(hy_eval.shape(), &[3, 1]);
123 assert_eq!(cy_eval.shape(), &[3, 1]);
124 }
125}