| { | |
| "format": "torch-dimensions/architecture", | |
| "io": { | |
| "cells_per_step": 48, | |
| "input": [ | |
| "B", | |
| "T", | |
| 6, | |
| 8, | |
| 1 | |
| ], | |
| "output": [ | |
| "B", | |
| "T", | |
| 6, | |
| 8, | |
| 32 | |
| ] | |
| }, | |
| "lattice": { | |
| "axes": [ | |
| { | |
| "dynamic": true, | |
| "name": "time", | |
| "size": null | |
| }, | |
| { | |
| "dynamic": false, | |
| "name": "h", | |
| "size": 6 | |
| }, | |
| { | |
| "dynamic": false, | |
| "name": "w", | |
| "size": 8 | |
| } | |
| ], | |
| "cells": { | |
| "dense": false, | |
| "present": 38, | |
| "present_rle": [ | |
| 0, | |
| 1, | |
| 1, | |
| 5, | |
| 1, | |
| 9, | |
| 1, | |
| 3, | |
| 1, | |
| 1, | |
| 1, | |
| 2, | |
| 1, | |
| 4, | |
| 1, | |
| 5, | |
| 1, | |
| 8, | |
| 2 | |
| ], | |
| "total": 48 | |
| }, | |
| "n_axes": 3, | |
| "names": [ | |
| "time", | |
| "h", | |
| "w" | |
| ], | |
| "rank": 2, | |
| "shape": [ | |
| 6, | |
| 8 | |
| ], | |
| "time": true | |
| }, | |
| "layers": [ | |
| { | |
| "axes": [ | |
| "time" | |
| ], | |
| "axis": "time", | |
| "axis_index": 0, | |
| "kind": "scan", | |
| "layer": 0, | |
| "mixer": "LSTMMixer", | |
| "n_params": 8448, | |
| "reverse": false | |
| }, | |
| { | |
| "axes": [ | |
| "h" | |
| ], | |
| "axis": "h", | |
| "axis_index": 1, | |
| "kind": "scan", | |
| "layer": 1, | |
| "mixer": "LSTMMixer", | |
| "n_params": 8448, | |
| "reverse": false | |
| }, | |
| { | |
| "axes": [ | |
| "w" | |
| ], | |
| "axis": "w", | |
| "axis_index": 2, | |
| "kind": "scan", | |
| "layer": 2, | |
| "mixer": "LSTMMixer", | |
| "n_params": 8448, | |
| "reverse": false | |
| }, | |
| { | |
| "axes": [ | |
| "time" | |
| ], | |
| "axis": "time", | |
| "axis_index": 0, | |
| "kind": "scan", | |
| "layer": 3, | |
| "mixer": "LSTMMixer", | |
| "n_params": 8448, | |
| "reverse": false | |
| }, | |
| { | |
| "axes": [ | |
| "h" | |
| ], | |
| "axis": "h", | |
| "axis_index": 1, | |
| "kind": "scan", | |
| "layer": 4, | |
| "mixer": "LSTMMixer", | |
| "n_params": 8448, | |
| "reverse": true | |
| }, | |
| { | |
| "axes": [ | |
| "w" | |
| ], | |
| "axis": "w", | |
| "axis_index": 2, | |
| "kind": "scan", | |
| "layer": 5, | |
| "mixer": "LSTMMixer", | |
| "n_params": 8448, | |
| "reverse": true | |
| } | |
| ], | |
| "model": { | |
| "d_input": 1, | |
| "d_model": 32, | |
| "kind": "LSTM", | |
| "n_layers": 6, | |
| "n_params": 51136 | |
| }, | |
| "nd_method": { | |
| "family": "scan", | |
| "name": "AxialScan" | |
| }, | |
| "sweeps": { | |
| "contracted_axes": [], | |
| "coverage": { | |
| "axes": [ | |
| { | |
| "backward": 0, | |
| "direction": "forward", | |
| "forward": 2, | |
| "index": 0, | |
| "layers": [ | |
| 0, | |
| 3 | |
| ], | |
| "name": "time" | |
| }, | |
| { | |
| "backward": 1, | |
| "direction": "both", | |
| "forward": 1, | |
| "index": 1, | |
| "layers": [ | |
| 1, | |
| 4 | |
| ], | |
| "name": "h" | |
| }, | |
| { | |
| "backward": 1, | |
| "direction": "both", | |
| "forward": 1, | |
| "index": 2, | |
| "layers": [ | |
| 2, | |
| 5 | |
| ], | |
| "name": "w" | |
| } | |
| ], | |
| "n_layers": 6, | |
| "pinned": [ | |
| "time" | |
| ], | |
| "unswept": [] | |
| }, | |
| "directions": { | |
| "h": "both", | |
| "time": "forward", | |
| "w": "both" | |
| }, | |
| "pinned_axes": [ | |
| "time" | |
| ], | |
| "unswept_axes": [] | |
| }, | |
| "version": 2 | |
| } | |