Skip to content

Instantly share code, notes, and snippets.

@dvruette
Last active June 11, 2024 10:49
Show Gist options
  • Save dvruette/f24fe0bcfa4fbb38773211733f27e3d5 to your computer and use it in GitHub Desktop.
Save dvruette/f24fe0bcfa4fbb38773211733f27e3d5 to your computer and use it in GitHub Desktop.
ada_numpy.py
from contextlib import contextmanager
@contextmanager
def set_numpy(name):
if name == "numpy":
import numpy as np
yield np
elif name == "jax":
import jax.numpy as jnp
yield jnp
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment