Remove jax enums

navix/_version.py

@@ -18,5 +18,5 @@
 # under the License.
 
 
-__version__ = "0.3.12"
+__version__ = "0.3.13"
 __version_info__ = tuple(int(i) for i in __version__.split(".") if i.isdigit())

navix/actions.py

@@ -26,10 +26,9 @@
 from .entities import Entities, State
 from .components import DISCARD_PILE_COORDS
 from .grid import translate, rotate, positions_equal
-from jax_enums import Enumerable as Enum
 
 
-class Directions(Enum):
+class Directions:
     EAST = 0
     SOUTH = 1
     WEST = 2

navix/entities.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 import dataclasses
-from typing import Dict, Tuple, Type, TypeVar
+from typing import Dict, Tuple, TypeVar
 
 import jax
 from jax import Array
 import jax.numpy as jnp
 from flax import struct
 from jax.random import KeyArray
-from jax_enums import Enumerable
 
 from .components import (
     Positionable,

pyproject.toml

@@ -9,7 +9,7 @@ requires = ["setuptools >= 50", "setuptools-scm[toml]>=6.2", "wheel"]
 name = "Navix"
 dynamic = ["version", "dependencies"]
 description = "Accelerated gridworld navigation with JAX for deep reinforcement learning"
-requires-python = ">=3.10"
+requires-python = ">=3.8"
 readme = "README.md"
 license = {file = "LICENSE", name = "Apache-2.0"}
 authors = [

requirements.txt

@@ -14,4 +14,3 @@ jax
 chex
 flax
 etils
-jax_enums