fix: registry

instadeepai · Sep 4, 2023 · d56588d · d56588d
1 parent 8931984
commit d56588d
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/matrax/__init__.py b/matrax/__init__.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import jax.numpy as jnp
 from jumanji.registration import make, register
 from jumanji.version import __version__
 
@@ -34,7 +33,7 @@
     )
     register(
         id=f"Penalty-{penalty_value}-stateful-v0",
-        entry_point="matrixgames.games:MatrixGame",
+        entry_point="matrax:MatrixGame",
         kwargs={
             "payoff_matrix": payoff_matrix,
             "keep_state": True,
@@ -53,7 +52,7 @@
 )
 register(
     id="Climbing-stateful-v0",
-    entry_point="matrixgames.games:MatrixGame",
+    entry_point="matrax:MatrixGame",
     kwargs={
         "payoff_matrix": climbing_game,
         "keep_state": True,
@@ -73,7 +72,7 @@
     )
     register(
         id=f"NoConflict-{_id}-stateful-v0",
-        entry_point="matrixgames.games:MatrixGame",
+        entry_point="matrax:MatrixGame",
         kwargs={
             "payoff_matrix": payoff_matrix,
             "keep_state": True,
@@ -93,7 +92,7 @@
     )
     register(
         f"Conflict-{_id}-stateful-v0",
-        entry_point="matrixgames.games:MatrixGame",
+        entry_point="matrix:MatrixGame",
         kwargs={
             "payoff_matrix": payoff_matrix,
             "keep_state": True,