Accept object for add method of ComputationalGraph class (#902)

Author: ishii-norimi

lib/model/neuralnetwork.js

@@ -1,6 +1,5 @@
 import Matrix from '../util/matrix.js'
 import Tensor from '../util/tensor.js'
-import Layer from './nns/layer/base.js'
 export { default as Layer } from './nns/layer/base.js'
 
 import ComputationalGraph from './nns/graph.js'
@@ -61,13 +60,11 @@ export default class NeuralNetwork {
 		}
 		const graph = new ComputationalGraph()
 		for (const cn of const_numbers) {
-			const cl = Layer.fromObject({ type: 'const', value: [[cn]] })
-			graph.add(cl, `__const_number_${cn}`, [])
+			graph.add({ type: 'const', value: [[cn]] }, `__const_number_${cn}`, [])
 		}
 
 		for (const l of layers) {
-			const cl = Layer.fromObject(l)
-			graph.add(cl, l.name, l.input)
+			graph.add(l, l.name, l.input)
 		}
 
 		return new NeuralNetwork(graph, optimizer)

lib/model/nns/graph.js

@@ -5,7 +5,11 @@ import { InputLayer, OutputLayer } from './layer/index.js'
 import ONNXImporter from './onnx/onnx_importer.js'
 
 /**
- * @typedef {import("./layer/index").PlainLayerObject & {input?: string | string[], name?: string}} LayerObject
+ * @ignore
+ * @typedef {import("./layer/index").PlainLayerObject} PlainLayerObject
+ */
+/**
+ * @typedef {PlainLayerObject & {input?: string | string[], name?: string}} LayerObject
  * @typedef {object} Node
  * @property {Layer} layer Layer
  * @property {string} name Name of the node
@@ -127,11 +131,14 @@ export default class ComputationalGraph {
 
 	/**
 	 * Add a layer.
-	 * @param {Layer} layer Added layer
+	 * @param {Layer | PlainLayerObject} layer Added layer
 	 * @param {string} [name] Node name
 	 * @param {string[] | string} [inputs] Input node names for the added layer
 	 */
 	add(layer, name, inputs = undefined) {
+		if (!(layer instanceof Layer)) {
+			layer = Layer.fromObject(layer)
+		}
 		let parentinfos = []
 		if (!inputs) {
 			if (this._nodes.length > 0) {

tests/lib/model/nns/graph.test.js

@@ -175,6 +175,16 @@ describe('Computational Graph', () => {
 			expect(graph.nodes[1].parents[0].subscript).toBeNull()
 		})
 
+		test('object', () => {
+			const graph = new ComputationalGraph()
+			graph.add({ type: 'input' })
+			graph.add({ type: 'tanh' })
+
+			expect(graph.nodes[1].parents).toHaveLength(1)
+			expect(graph.nodes[1].parents[0].index).toBe(0)
+			expect(graph.nodes[1].parents[0].subscript).toBeNull()
+		})
+
 		test('string input', () => {
 			const graph = new ComputationalGraph()
 			graph.add(Layer.fromObject({ type: 'input' }), 'in')