readme.md

README.md

+
+
+# 下载jax和jaxlib
+
+```
+wget https://cancon.hpccube.com:65024/directlink/4/jax/DAS1.1/jaxlib-0.4.23+das1.1.git387bd43.abi1.dtk2404-cp39-cp39-manylinux_2_31_x86_64.whl
+wget https://cancon.hpccube.com:65024/directlink/4/jax/DAS1.1/jax-0.4.23+das1.1.git387bd43.abi1.dtk2404-py3-none-any.whl
+
+```
+
+
+# conda 环境
+
+```
+conda create -n pymc3 python=3.9
+conda activate pymc3
+pip install jax-0.4.23+das1.1.git387bd43.abi1.dtk2404-py3-none-any.whl
+pip install jaxlib-0.4.23+das1.1.git387bd43.abi1.dtk2404-cp39-cp39-manylinux_2_31_x86_64.whl
+pip install  pymc==5.9.1
+pip install numpyro==0.14.0
+pip install seaborn==0.13.2
+pip install scipy==1.12.0 -i https://pypi.tuna.tsinghua.edu.cn/simple 
+```
+
+
+# 测试jax是否在gpu上可行：
+  ```
+import jax
+import pymc3 as pm
+jax.default_backend()
+jax.devices()
+```
+
+# 测试采样时间，GPU负载等情况
+
+```
+import jax
+import pymc as pm
+import numpy as np
+import pytensor as pt
+pt.config.floatX = "float32"
+np.random.seed(123) 
+n =10000
+X = np.random.randn(n) 
+Y =3* X + np.random.randn(n)
+# 定义PyMC3模型
+with pm.Model() as model:
+    alpha = pm.Normal('alpha', mu=0, sigma=10)     
+    beta = pm.Normal('beta', mu=0, sigma=10)     
+    sigma = pm.HalfNormal('sigma', sigma=1)
+    mu = alpha + beta * X
+    Y_obs = pm.Normal('Y_obs', mu=mu, sigma=sigma, observed=Y)
+    trace = pm.sample(1000, nuts_sampler="numpyro",return_inferencedata=False)
+print(pm.summary(trace))
+  
+
+```
+
+   
\ No newline at end of file

test_gpu.py

+import jax
+import pymc as pm
+print(jax.default_backend())
+print(jax.devices())
\ No newline at end of file

test_sample_time.py

+import jax
+import pymc as pm
+import numpy as np
+import pytensor as pt
+pt.config.floatX = "float32"
+np.random.seed(123) 
+n =10000
+X = np.random.randn(n) 
+Y =3* X + np.random.randn(n)
+# 定义PyMC3模型
+with pm.Model() as model:
+    alpha = pm.Normal('alpha', mu=0, sigma=10)     
+    beta = pm.Normal('beta', mu=0, sigma=10)     
+    sigma = pm.HalfNormal('sigma', sigma=1)
+    mu = alpha + beta * X
+    Y_obs = pm.Normal('Y_obs', mu=mu, sigma=sigma, observed=Y)
+    trace = pm.sample(1000, nuts_sampler="numpyro",return_inferencedata=False)
+print(pm.summary(trace))
\ No newline at end of file