Merged files, removed pmap_experimental.

graphcore-research / pyscf-ipu

PySCF on IPU

https://github.com/graphcore-research/pyscf-ipu#readme

Apache License 2.0

41 stars 3 forks source link

Merged files, removed pmap_experimental. #97

Closed AlexanderMath closed 10 months ago

AlexanderMath commented 10 months ago

Removing pmap_experimental_nanoDFT by merging changes into nanoDFT. All the jax.pmap stuff becomes identity when ndevices=1.

AlexanderMath commented 10 months ago

@hatemhelal IPU_MODEL fails when using int16/float16 for the indices. Can we do something like if ipu_model: dense_ERI=True?

hatemhelal commented 10 months ago

@hatemhelal IPU_MODEL fails when using int16/float16 for the indices. Can we do something like if ipu_model: dense_ERI=True?

yes I think the way to do it would be:

pyscf_ipu.experimental.device.has_ipu() and "JAX_IPU_USE_MODEL" in os.environ

It is handy to have utilities for feature testing if we are running on an IPU and if that IPU is a model so seems worthwhile to put that in a common place.

review-notebook-app[bot] commented 10 months ago

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.