Add Bohrium support

[jamesavery]

TODO:

• [x] imed.py (turn imed_loss back on in test_run_kuro.py and test_run_mackey_2d.py
• [ ] make it usable with params.backen = 'bohrium'
• [ ] bohrium 2D convolutions?

[nmheim]

awesome seems to be working! @jamesavery should we do our own 2D convolution in bohrium?

[nmheim]

Comparing pytest -k mackey_2d for bohrium and numpy: NUMPY:

        5.631 predict
            4.853 forward
                0.141 tanh
                1.996 state_map
                2.700 input_map
                    0.053 gradient
                    0.062 dct
                    2.439 conv
                    0.122 pixels
       20.603 optimize
            0.408 IMED matmul on wout
           18.815 pseudo_inverse_svd
               17.812 svd
            0.737 IMED matmul on inputs and labels
            0.052 IMED prepare
            0.504 IMED initialize
       23.916 forward_states_only
           23.792 forward
                0.631 tanh
               10.249 state_map
               12.839 input_map
                    0.268 gradient
                    0.302 dct
                   11.631 conv
                    0.530 pixels

BOHRIUM

        4.772 predict                                                                                                                                                                                                                                                                     
            3.680 forward                                                                                                                                                                                                                                                                 
                0.132 tanh                                                                                                                                                                                                                                                                
                0.470 state_map                                                                                                                                                                                                                                                           
                2.415 input_map                                                                                                                                                                                                                                                           
                    0.675 gradient                                                                                                                                                                                                                                                        
                    0.056 dct                                                                                                                                                                                                                                                             
                    1.579 conv                                                                                                                                                                                                                                                            
                    0.082 pixels                                                                                                                                                                                                                                                          
       23.538 optimize                                                                                                                                                                                                                                                                    
            0.316 IMED matmul on wout                                                                                                                                                                                                                                                     
           21.874 pseudo_inverse_svd                                                                                                                                                                                                                                                      
               17.426 svd                                                                                                                                                                                                                                                                 
            0.764 IMED matmul on inputs and labels                                                                                                                                                                                                                                        
            0.052 IMED prepare                                                                                                                                                                                                                                                            
            0.527 IMED initialize                                                                                                                                                                                                                                                         
       35.976 forward_states_only                                                                                                                                                                                                                                                         
           35.506 forward                                                                                                                                                                                                                                                                 
                1.309 tanh                                                                                                                                                                                                                                                                
                3.617 state_map                                                                                                                                                                                                                                                           
               24.054 input_map                                                                                                                                                                                                                                                           
                    6.702 gradient                                                                                                                                                                                                                                                        
                    0.549 dct                                                                                                                                                                                                                                                             
                   15.795 conv                                                                                                                                                                                                                                                            
                    0.784 pixels

[jamesavery]

Something funny is causing massive slowdown for Bohrium - need to talk to Mads when he's back from holiday.