1 files changed, 21 insertions, 14 deletions

diff --git a/src/lib/src/neuralnet_impl.h b/src/lib/src/neuralnet_impl.h
index f5a9c63..935c5ea 100644
--- a/src/lib/src/neuralnet_impl.h
+++ b/src/lib/src/neuralnet_impl.h
@@ -2,22 +2,29 @@
 
 #include <neuralnet/matrix.h>
 
+#include <stdbool.h>
+
+/// Linear layer parameters.
+typedef struct nnLinearImpl {
+  nnMatrix weights;
+  nnMatrix biases;
+  bool     owned; /// Whether the library owns the weights and biases.
+} nnLinearImpl;
+
+/// Neural network layer.
+typedef struct nnLayerImpl {
+  nnLayerType type;
+  int         input_size;
+  int         output_size;
+  union {
+    nnLinearImpl linear;
+  };
+} nnLayerImpl;
+
 /// Neural network object.
-///
-/// We store the transposes of the weight matrices so that we can do forward
-/// passes with a minimal amount of work. That is, if in paper we write:
-///
-///   [w11 w21]
-///   [w12 w22]
-///
-/// then the weight matrix in memory is stored as the following array:
-///
-///   w11 w12 w21 w22
 typedef struct nnNeuralNetwork {
-  int           num_layers;  // Number of non-input layers (hidden + output).
-  nnMatrix*     weights;     // One matrix per non-input layer.
-  nnMatrix*     biases;      // One vector per non-input layer.
-  nnActivation* activations; // One per non-input layer.
+  int          num_layers; // Number of non-input layers (hidden + output).
+  nnLayerImpl* layers;     // One per non-input layer.
 } nnNeuralNetwork;
 
 /// A query object that holds all the memory necessary to query a network.