1 files changed, 32 insertions, 19 deletions

diff --git a/src/lib/include/neuralnet/neuralnet.h b/src/lib/include/neuralnet/neuralnet.h
index 05c9406..f122c2a 100644
--- a/src/lib/include/neuralnet/neuralnet.h
+++ b/src/lib/include/neuralnet/neuralnet.h
@@ -1,32 +1,45 @@
 #pragma once
 
+#include <neuralnet/matrix.h>
 #include <neuralnet/types.h>
 
-typedef struct nnMatrix nnMatrix;
-
 typedef struct nnNeuralNetwork nnNeuralNetwork;
 typedef struct nnQueryObject   nnQueryObject;
 
-/// Neuron activation.
-typedef enum nnActivation {
-  nnIdentity,
+/// Linear layer parameters.
+///
+/// Either one of the following must be set:
+///   a) Training:  input and output sizes.
+///   b) Inference: weights + biases.
+typedef struct nnLinearParams {
+  int      input_size;
+  int      output_size;
+  nnMatrix weights;
+  nnMatrix biases;
+} nnLinearParams;
+
+/// Layer type.
+typedef enum nnLayerType {
+  nnLinear,
   nnSigmoid,
   nnRelu,
-} nnActivation;
+} nnLayerType;
+
+/// Neural network layer.
+typedef struct nnLayer {
+  nnLayerType type;
+  union {
+    nnLinearParams linear;
+  };
+} nnLayer;
 
 /// Create a network.
 nnNeuralNetwork* nnMakeNet(
-    int num_layers, const int* layer_sizes, const nnActivation* activations);
+    const nnLayer* layers, int num_layers, int input_size);
 
 /// Delete the network and free its internal memory.
 void nnDeleteNet(nnNeuralNetwork**);
 
-/// Set the network's weights.
-void nnSetWeights(nnNeuralNetwork*, const R* weights);
-
-/// Set the network's biases.
-void nnSetBiases(nnNeuralNetwork*, const R* biases);
-
 /// Query the network.
 ///
 /// |input| is a matrix of inputs, one row per input and as many columns as the
@@ -42,10 +55,10 @@ void nnQueryArray(
 
 /// Create a query object.
 ///
-/// The query object holds all the internal memory required to query a network.
-/// Query objects allocate all memory up front so that network queries can run
-/// without additional memory allocation.
-nnQueryObject* nnMakeQueryObject(const nnNeuralNetwork*, int num_inputs);
+/// The query object holds all the internal memory required to query a network
+/// with batches of the given size. Memory is allocated up front so that network
+/// queries can run without additional memory allocation.
+nnQueryObject* nnMakeQueryObject(const nnNeuralNetwork*, int batch_size);
 
 /// Delete the query object and free its internal memory.
 void nnDeleteQueryObject(nnQueryObject**);
@@ -60,7 +73,7 @@ int nnNetInputSize(const nnNeuralNetwork*);
 int nnNetOutputSize(const nnNeuralNetwork*);
 
 /// Return the layer's input size.
-int nnLayerInputSize(const nnMatrix* weights);
+int nnLayerInputSize(const nnNeuralNetwork*, int layer);
 
 /// Return the layer's output size.
-int nnLayerOutputSize(const nnMatrix* weights);
+int nnLayerOutputSize(const nnNeuralNetwork*, int layer);