{
 "format": "torch-dimensions/architecture",
 "io": {
  "cells_per_step": 240,
  "input": [
   "B",
   "T",
   3,
   4,
   5,
   4,
   1
  ],
  "output": [
   "B",
   "T",
   3,
   4,
   5,
   4,
   32
  ]
 },
 "lattice": {
  "axes": [
   {
    "dynamic": true,
    "name": "time",
    "size": null
   },
   {
    "dynamic": false,
    "name": "depth",
    "size": 3
   },
   {
    "dynamic": false,
    "name": "row",
    "size": 4
   },
   {
    "dynamic": false,
    "name": "col",
    "size": 5
   },
   {
    "dynamic": false,
    "name": "group",
    "size": 4
   }
  ],
  "cells": {
   "dense": false,
   "present": 169,
   "present_rle": [
    0,
    6,
    2,
    1,
    1,
    1,
    1,
    5,
    3,
    2,
    1,
    8,
    1,
    12,
    1,
    4,
    2,
    2,
    2,
    2,
    2,
    1,
    1,
    7,
    2,
    1,
    1,
    2,
    1,
    2,
    1,
    1,
    1,
    4,
    1,
    8,
    1,
    3,
    1,
    3,
    1,
    6,
    1,
    1,
    2,
    4,
    2,
    1,
    1,
    10,
    1,
    1,
    1,
    3,
    2,
    10,
    2,
    1,
    1,
    1,
    2,
    1,
    1,
    6,
    1,
    2,
    1,
    8,
    1,
    1,
    2,
    1,
    1,
    1,
    2,
    3,
    1,
    2,
    1,
    4,
    1,
    3,
    1,
    1,
    4,
    1,
    1,
    1,
    1,
    3,
    1,
    6,
    1,
    3,
    2,
    3,
    1,
    2,
    3,
    1,
    1,
    3,
    1
   ],
   "total": 240
  },
  "n_axes": 5,
  "names": [
   "time",
   "depth",
   "row",
   "col",
   "group"
  ],
  "rank": 4,
  "shape": [
   3,
   4,
   5,
   4
  ],
  "time": true
 },
 "layers": [
  {
   "axes": [
    "time"
   ],
   "axis": "time",
   "axis_index": 0,
   "kind": "scan",
   "layer": 0,
   "mixer": "S4DMixer",
   "n_params": 4224,
   "reverse": false
  },
  {
   "axes": [
    "depth"
   ],
   "axis": "depth",
   "axis_index": 1,
   "kind": "scan",
   "layer": 1,
   "mixer": "S4DMixer",
   "n_params": 4224,
   "reverse": false
  },
  {
   "axes": [
    "row"
   ],
   "axis": "row",
   "axis_index": 2,
   "kind": "scan",
   "layer": 2,
   "mixer": "S4DMixer",
   "n_params": 4224,
   "reverse": false
  },
  {
   "axes": [
    "col"
   ],
   "axis": "col",
   "axis_index": 3,
   "kind": "scan",
   "layer": 3,
   "mixer": "S4DMixer",
   "n_params": 4224,
   "reverse": false
  },
  {
   "axes": [
    "group"
   ],
   "axis": "group",
   "axis_index": 4,
   "kind": "scan",
   "layer": 4,
   "mixer": "S4DMixer",
   "n_params": 4224,
   "reverse": false
  },
  {
   "axes": [
    "time"
   ],
   "axis": "time",
   "axis_index": 0,
   "kind": "scan",
   "layer": 5,
   "mixer": "S4DMixer",
   "n_params": 4224,
   "reverse": false
  },
  {
   "axes": [
    "depth"
   ],
   "axis": "depth",
   "axis_index": 1,
   "kind": "scan",
   "layer": 6,
   "mixer": "S4DMixer",
   "n_params": 4224,
   "reverse": false
  },
  {
   "axes": [
    "row"
   ],
   "axis": "row",
   "axis_index": 2,
   "kind": "scan",
   "layer": 7,
   "mixer": "S4DMixer",
   "n_params": 4224,
   "reverse": false
  }
 ],
 "model": {
  "d_input": 1,
  "d_model": 32,
  "kind": "S4D",
  "n_layers": 8,
  "n_params": 34368
 },
 "nd_method": {
  "family": "scan",
  "name": "AxialScan"
 },
 "sweeps": {
  "contracted_axes": [],
  "coverage": {
   "axes": [
    {
     "backward": 0,
     "direction": "forward",
     "forward": 2,
     "index": 0,
     "layers": [
      0,
      5
     ],
     "name": "time"
    },
    {
     "backward": 0,
     "direction": "forward",
     "forward": 2,
     "index": 1,
     "layers": [
      1,
      6
     ],
     "name": "depth"
    },
    {
     "backward": 0,
     "direction": "forward",
     "forward": 2,
     "index": 2,
     "layers": [
      2,
      7
     ],
     "name": "row"
    },
    {
     "backward": 0,
     "direction": "forward",
     "forward": 1,
     "index": 3,
     "layers": [
      3
     ],
     "name": "col"
    },
    {
     "backward": 0,
     "direction": "forward",
     "forward": 1,
     "index": 4,
     "layers": [
      4
     ],
     "name": "group"
    }
   ],
   "n_layers": 8,
   "pinned": [
    "time",
    "depth",
    "row",
    "col",
    "group"
   ],
   "unswept": []
  },
  "directions": {
   "col": "forward",
   "depth": "forward",
   "group": "forward",
   "row": "forward",
   "time": "forward"
  },
  "pinned_axes": [
   "time",
   "depth",
   "row",
   "col",
   "group"
  ],
  "unswept_axes": []
 },
 "version": 2
}