Fix linter errors.

Author: aterzis-google

objax/nn/layers.py

@@ -331,6 +331,7 @@ def __call__(self, x: JaxArray) -> JaxArray:
         self.avg.value += (self.avg.value - x) * (self.momentum - 1)
         return self.avg.value
 
+
 class SimpleRNN(Module):
     """Simple Recurrent Neural Network (RNN) block."""
 
@@ -363,12 +364,10 @@ def __init__(self,
         self.w_hh = TrainVar(w_init((self.nstate, self.nstate)))
         self.b_h = TrainVar(jn.zeros(self.nstate))
 
+        self.output_layer = Linear(self.nstate, self.num_outputs, w_init=w_init)
 
-        self.output_layer = Linear(self.nstate, self.num_outputs, w_init = w_init)
-
-    def __call__(self, inputs: JaxArray,
-                 initial_state: JaxArray = None,
-                 only_return_final = False) -> Tuple[JaxArray, JaxArray]:
+    def __call__(self, inputs: JaxArray, initial_state: JaxArray = None,
+                 only_return_final: bool = False) -> Tuple[JaxArray, JaxArray]:
         """Forward pass through RNN.
 
         Args:
@@ -383,7 +382,7 @@ def __call__(self, inputs: JaxArray,
         """
         outputs = []
 
-        if initial_state == None:
+        if initial_state is None:
             state = jn.zeros((inputs.shape[0], self.nstate))
         else:
             state = initial_state

tests/simple_rnn.py

@@ -16,24 +16,21 @@
 
 import unittest
 
-import numpy as np
 import jax.numpy as jn
-import tensorflow as tf
 
-import objax
 from objax.nn.layers import SimpleRNN
 from objax.functional import one_hot
 from objax.functional.core.activation import relu
 from objax.nn.init import identity
-from objax.zoo.resnet_v2 import convert_keras_model, load_pretrained_weights_from_keras
+
 
 class TestSimpleRNN(unittest.TestCase):
 
     def test_simple_rnn(self):
         nin = nout = 3
         batch_size = 1
         num_hiddens = 1
-        model = SimpleRNN(num_hiddens, nin, nout, activation = relu, w_init = identity)
+        model = SimpleRNN(num_hiddens, nin, nout, activation=relu, w_init=identity)
 
         X = jn.arange(batch_size)
         X_one_hot = one_hot(X, nin)
@@ -48,5 +45,6 @@ def test_simple_rnn(self):
         Z, _ = model(X_one_hot, state)
         self.assertTrue(jn.array_equal(Z, jn.array([[3., 0., 0.]])))
 
+
 if __name__ == '__main__':
     unittest.main()