llandsmeer

Spike time function

@jax.custom_jvp
def tpre(t, v, vdot):
    return t

@tpre.defjvp
def tpre_jvp(primals, tangents):
    t, v, vdot = primals

llandsmeer

# pip install keystone-engine

import inspect
import struct
from keystone import Ks, KS_ARCH_X86, KS_MODE_64
import ctypes
import mmap

def map_type(x):
    if isinstance(x, int): return 'i64'

llandsmeer

npx esbuild main.tsx --bundle --format=esm --outfile=dist/main.js --jsx=automatic --jsx-import-source=preact

llandsmeer

import jax
import tqdm
from jax import numpy as jnp
from brax.io import mjcf
from brax.envs import PipelineEnv, State
from brax.io import html
from brax.generalized import pipeline


sys = mjcf.load('osc-mouse.xml')

llandsmeer

set -ex

. ../env/bin/activate

python3 jax_test.py

iree-opt ./jax.stablehlo.mlir \
    --iree-stablehlo-input-transformation-pipeline  \
    --convert-scf-to-cf \
    > jax.linalg.mlir

llandsmeer

Build

git clone https://github.com/UCLA-VAST/Stream-HLS
cd Stream-HLS
sudo add-apt-repository ppa:deadsnakes/ppa
sudo apt update
sudo apt install python3.11 python3.11-venv
python3.11 -m venv env
. env/bin/activate

llandsmeer

import jax
import jax.numpy as jnp

import matplotlib.pyplot as plt

def normalize(x):
    return x / jnp.linalg.norm(x)

def vec(x, y, z):
    return jnp.array([x, y, z], dtype='float32')

llandsmeer

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import functools
import typing

def main():
    params = IOCellParams(g_CaL=1.2)
    state0 = IOCellState.make(params)
    @jax.jit

llandsmeer

for model in paper_experiments/nml_standalones/*/
do
    echo -n "$(basename "${model}")"
    for fun in log exp sqrt sin cos tan sinh cosh tanh abs ceil factorial random
    do
        if rg \
            -q \
            -i \
            -g '*.nml' \
            '\b'"${fun}"'\b' \

llandsmeer

'Conv'                         #############################################################################################################################################################
'MaxPool'                      #################################################################################################################################################
'Relu'                         #############################################################################################################################################
'Reshape'                      ##################################################################################################################################
'Softmax'                      ##########################################################################################################
'Concat'                       ##################################################################################################
'Add'                          ########################