Skip to content

Instantly share code, notes, and snippets.

View kazewong's full-sized avatar

Kaze Wong kazewong

95 followers · 12 following
View GitHub Profile
@kazewong
kazewong / IntroductionToFlax.ipynb
Last active June 29, 2022 18:22
Introduction to flax
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
@kazewong
kazewong / GW150914.ipynb
Created October 6, 2022 03:05
Notebook for running GW150914 with FlowMC
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
@kazewong
kazewong / spline.py
Last active December 2, 2023 16:55
Cubic Spline with lineax
import lineax as lx
import jax.numpy as jnp
import jax
from jaxtyping import Float, Array
class CubicSpline:
x_grid: Float[Array, str("batch")] # input x data
y_grid: Float[Array, str("n")] # input y data
@kazewong
kazewong / test_shard_model.py
Created July 7, 2024 14:59
Sharding equinox linear model
import math
from typing import Any, Literal, Union
import numpy as np
import jax
import jax.numpy as jnp
from jax._src.distributed import initialize
from jaxtyping import PRNGKeyArray
import equinox as eqx
from equinox._misc import default_floating_dtype