chore: undo ruff linting

instadeepai · Jan 14, 2024 · 1a122f4 · 1a122f4
1 parent 0cf3550
commit 1a122f4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/matrax/__init__.py b/matrax/__init__.py
@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from jumanji.registration import register
+from jumanji.registration import make, register
+from jumanji.version import __version__
 
+from matrax.env import MatrixGame
 from matrax.games import climbing_game, conflict_games, no_conflict_games, penalty_games
+from matrax.types import Observation, State
 
 """Environment Registration"""
 

diff --git a/matrax/env_test.py b/matrax/env_test.py
@@ -21,9 +21,8 @@
 from jumanji.testing.pytrees import assert_is_jax_array_tree
 from jumanji.types import TimeStep
 
-from matrax.env import MatrixGame
+from matrax import MatrixGame, State
 from matrax.games import climbing_game
-from matrax.types import State
 
 
 @pytest.fixture