Skip to content

Improve device information in `netket.tools.info`

Vicentini Filippo requested to merge tools.info into master

Created by: femtobit

This fixes as small display bug in netket.tools.info:

Before:

$ python3 -m netket.tools.info
[...]
# MPI4JAX
  - HAS_GPU_EXT      : <function has_cuda_support at 0x...>
[...]

After:

$ python3 -m netket.tools.info
[...]
# MPI4JAX
  - HAS_GPU_EXT      : True
[...]

Edit: Since I already had that file open: I've added some code to print nicer looking device information:

# Jax 
  - backends         : ['cpu', 'gpu']
    - cpu            : ['<0: cpu>']
    - gpu            : ['<0: Tesla V100-PCIE-32GB>', '<1: Tesla V100-PCIE-32GB>'

Merge request reports