torch-dimensions / viewer /src /samples /lstm_2d_sparse.json
Celsia's picture
Upload folder using huggingface_hub
ecc81b3 verified
Raw
History Blame Contribute Delete
2.75 kB
{
"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
}