[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

examples/rcn/helpers.py

@@ -18,10 +18,10 @@ def get_number_of_states(hps, vps):
 
 def initialize_evidences(inf_img, frcs, hps, vps, neg_inf=-1000):
     M = get_number_of_states(hps, vps)
-    
+
     preproc_layer = Preproc(cross_channel_pooling=True)
     bu_msg = preproc_layer.fwd_infer(inf_img)
-    
+
     evidence_updates = {}
 
     for idx in range(frcs.shape[0]):

examples/rcn/inference_pgmax_small.py

@@ -17,11 +17,12 @@
 # %matplotlib inline
 import os
 import time
-from jax import jit, tree_util
-from jax import numpy as jnp
+
 import matplotlib.pyplot as plt
 import numpy as np
-
+from jax import jit
+from jax import numpy as jnp
+from jax import tree_util
 from load_data import get_mnist_data_iters
 from preproc import Preproc
 
@@ -53,17 +54,16 @@
 # -
 
 
-
 # %% [markdown]
 # ## 2. Load the model
 #
 #
 
 # %%
 directory = f"/storage/users/skushagra/pgmax_rcn_artifacts/model_science_{train_size}_{hps}_{vps}"
-frcs = np.load(f"{directory}/frcs.npy", allow_pickle=True, encoding='latin1')
-edges = np.load(f"{directory}/edges.npy", allow_pickle=True, encoding='latin1')
-M = (2 * hps + 1) * (2 * vps + 1) + 1 
+frcs = np.load(f"{directory}/frcs.npy", allow_pickle=True, encoding="latin1")
+edges = np.load(f"{directory}/edges.npy", allow_pickle=True, encoding="latin1")
+M = (2 * hps + 1) * (2 * vps + 1) + 1
 
 # %% [markdown]
 # ## 3. Visualize loaded model.
@@ -88,7 +88,6 @@
 plt.imshow(img, cmap="gray")
 
 
-
 # %% [markdown]
 # ## 4. Make pgmax graph
 #
@@ -123,7 +122,7 @@
 # %%
 
 # %% [markdown]
-# ## 3.2.1 Pre-compute the valid configs for different perturb radii. 
+# ## 3.2.1 Pre-compute the valid configs for different perturb radii.
 #
 #
 
@@ -150,6 +149,7 @@ def valid_configs(r):
 
     return np.stack([rows, cols], axis=1)
 
+
 max_perturb_radii = 25
 phis = []
 for r in range(max_perturb_radii):
@@ -181,7 +181,6 @@ def valid_configs(r):
 print(f"Creating factors took {end-start:.3f} seconds.")
 
 
-
 # %% [markdown]
 # ## 4. Run inference
 #
@@ -196,7 +195,7 @@ def valid_configs(r):
 def initialize_evidences(test_img, frcs, hps, vps):
     preproc_layer = Preproc(cross_channel_pooling=True)
     bu_msg = preproc_layer.fwd_infer(test_img)
-    
+
     evidence_updates = {}
     for idx in range(frcs.shape[0]):
         frc = frcs[idx]
@@ -237,25 +236,25 @@ def initialize_evidences(test_img, frcs, hps, vps):
     print(f"Initializing evidences took {end-start:.3f} seconds for image {test_idx}.")
 
     start = end
-    map_states = graph.decode_map_states(get_beliefs_fn(run_bp_fn(evidence_updates=evidence_updates)))
+    map_states = graph.decode_map_states(
+        get_beliefs_fn(run_bp_fn(evidence_updates=evidence_updates))
+    )
     end = time.time()
     print(f"Max product inference took {end-start:.3f} seconds for image {test_idx}.")
 
     map_states_dict[test_idx] = map_states
     start = end
     score = tree_util.tree_multimap(
-            lambda evidence, map: jnp.sum(
-                evidence[jnp.arange(map.shape[0]), map]
-            ), 
-            evidence_updates, 
-            map_states
-        )
-    for ii in score: scores[test_idx, ii] = score[ii]
+        lambda evidence, map: jnp.sum(evidence[jnp.arange(map.shape[0]), map]),
+        evidence_updates,
+        map_states,
+    )
+    for ii in score:
+        scores[test_idx, ii] = score[ii]
     end = time.time()
     print(f"Computing scores took {end-start:.3f} seconds for image {test_idx}.")
 
-    #scores[test_idx, :] = score
-
+    # scores[test_idx, :] = score
 
 
 # %% [markdown]
@@ -270,7 +269,6 @@ def initialize_evidences(test_img, frcs, hps, vps):
 print(f"accuracy = {accuracy}")
 
 
-
 # %% [markdown]
 # ## 6. Visualize predictions (backtrace)
 #
@@ -281,12 +279,10 @@ def initialize_evidences(test_img, frcs, hps, vps):
 plt.imshow(test_set[test_idx][0], cmap="gray")
 
 
-
 # %%
 # ## 6.1 Backtrace of some models on this test image
 
 
-
 # %%
 # +
 map_states = map_states_dict[test_idx]
@@ -306,11 +302,10 @@ def initialize_evidences(test_img, frcs, hps, vps):
 plt.figure(figsize=(15, 15))
 
 for k, index in enumerate(range(0, len(train_set), 5)):
-    plt.subplot(1, 4, 1+k)
+    plt.subplot(1, 4, 1 + k)
     plt.title(f" Model {int(train_labels[index])}")
     plt.imshow(imgs[index, :, :], cmap="gray")
 # -
 
 
-
 # %%

examples/rcn/model_create/save_model.py

@@ -10,8 +10,8 @@
 import time
 
 import numpy as np
-from learning import Model
 from helpers import get_number_of_states, index_to_rc, rc_to_index
+from learning import Model
 from load_data import get_mnist_data_iters