Skip to content

Instantly share code, notes, and snippets.

@shoyer
shoyer / jax-einsum-primitive.ipynb
Created May 18, 2023 17:27
JAX einsum primitive .ipynb
View jax-einsum-primitive.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
@shoyer
shoyer / xarray_tensorstore.py
Created November 26, 2022 23:54
xarray zarr via tensor store
View xarray_tensorstore.py
# Copyright 2022 Google LLC.
# SPDX-License-Identifier: Apache-2.0
import tensorstore
import json
import os.path
import fsspec
import xarray
import xarray.backends
@shoyer
shoyer / jaxopt-implicit-optimization-gradients-with-jax-md.ipynb
Created February 8, 2022 01:30
JAXopt implicit optimization gradients with JAX-MD.ipynb
View jaxopt-implicit-optimization-gradients-with-jax-md.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
@shoyer
shoyer / exploring-kuramoto-sivashinsky-in-jax-cfd.ipynb
Last active October 14, 2021 23:52
Exploring Kuramoto-Sivashinsky in JAX-CFD.ipynb
View exploring-kuramoto-sivashinsky-in-jax-cfd.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
@shoyer
shoyer / jax-new-scan-benchmarking.ipynb
Last active August 29, 2021 23:28
JAX new scan benchmarking.ipynb
View jax-new-scan-benchmarking.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
@shoyer
shoyer / new-gmres-benchmark.ipynb
Created December 1, 2020 21:03
new gmres benchmark.ipynb
View new-gmres-benchmark.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
View fast_groupby_bins_agg.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
View dask-memory-bug.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
@shoyer
shoyer / copy-of-benchmark_gmres.ipynb
Last active September 16, 2020 19:03
Copy of benchmark_gmres.ipynb
View copy-of-benchmark_gmres.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
@shoyer
shoyer / jax-harmonic-oscillator-odeint.ipynb
Last active December 11, 2021 16:43
JAX harmonic oscillator odeint.ipynb
View jax-harmonic-oscillator-odeint.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.