From 9fab8179ef0cf1f93cbcfad07b31d3992c6269b1 Mon Sep 17 00:00:00 2001
From: Tianyi Hao <tianyihao@alumni.stanford.edu>
Date: Mon, 29 Jan 2024 02:35:52 -0600
Subject: [PATCH] alias jax.scipy #17

---
 autoray/autoray.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/autoray/autoray.py b/autoray/autoray.py
index dbb5617..cec7751 100644
--- a/autoray/autoray.py
+++ b/autoray/autoray.py
@@ -1495,6 +1495,7 @@ def jax_to_numpy(x):
 
 
 _BACKEND_ALIASES["jaxlib"] = "jax"
+_MODULE_ALIASES["jax.scipy"] = "jax.scipy"
 _MODULE_ALIASES["jax"] = "jax.numpy"
 _SUBMODULE_ALIASES["jax", "complex"] = "jax.lax"
 _SUBMODULE_ALIASES["jax", "linalg.expm"] = "jax.scipy.linalg"