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'                          ########################

llandsmeer

import jax
import iree.compiler as compiler

def YOUR_FUNCTION(x):
    return x + 1

input_sample = [
    jax.numpy.array([0.,])
]

llandsmeer

Excerpt from this issue: patrick-kidger/diffrax#202

llandsmeer

import torch
import numpy as np
import norse.torch as norse
import matplotlib.pyplot as plt
import tqdm.notebook as tqdm

_ = torch.manual_seed(0)

class MyModule(torch.nn.Module):
    def __init__(self, p=norse.LIFParameters()):