fix error: jax_array config option is removed in jax v0.4.14

renderer/pipeline.py

@@ -40,7 +40,8 @@
     ZBuffer,
 )
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 RowIndices = Integer[Array, "row_batches row_batch_size"]
 """Indices of the rows in the buffers to be processed in this batch."""

renderer/shader.py

@@ -28,7 +28,8 @@
     Vec4f,
 )
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 ID: TypeAlias = IntV
 

renderer/shaders/depth.py

@@ -14,7 +14,8 @@
 from ..shader import ID, PerVertex, Shader
 from ..types import Vec4f
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 
 class DepthExtraInput(NamedTuple):

renderer/shaders/gouraud.py

@@ -15,7 +15,8 @@
 from ..shader import ID, PerFragment, PerVertex, Shader
 from ..types import BoolV, Colour, FloatV, LightSource, Vec2f, Vec3f, Vec4f
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 
 class GouraudExtraInput(NamedTuple):

renderer/shaders/gouraud_texture.py

@@ -16,7 +16,8 @@
 from ..shader import ID, MixerOutput, PerFragment, PerVertex, Shader
 from ..types import BoolV, Colour, FloatV, LightSource, Texture, Vec2f, Vec3f, Vec4f
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 
 class GouraudTextureExtraInput(NamedTuple):

renderer/shaders/phong.py

@@ -16,7 +16,8 @@
 from ..shader import ID, MixerOutput, PerFragment, PerVertex, Shader
 from ..types import BoolV, Colour, LightSource, Texture, Vec2f, Vec3f, Vec4f
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 
 class PhongTextureExtraInput(NamedTuple):

renderer/shaders/phong_darboux.py

@@ -35,7 +35,8 @@
     Vec4f,
 )
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 Triangle3f: TypeAlias = Float[Array, "3 3"]
 Triangle2f: TypeAlias = Float[Array, "3 2"]

renderer/shaders/phong_reflection.py

@@ -29,7 +29,8 @@
     Vec4f,
 )
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 
 class PhongReflectionTextureExtraInput(NamedTuple):

renderer/shaders/phong_reflection_shadow.py

@@ -30,7 +30,8 @@
     Vec4f,
 )
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 
 class PhongReflectionShadowTextureExtraInput(NamedTuple):

renderer/types.py

@@ -40,7 +40,8 @@
     "Buffers",
 ]
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 BoolV: TypeAlias = Bool[Array, ""]
 """JAX Array with single bool value.""" ""

test_resources/utils.py

@@ -11,7 +11,8 @@
 from renderer import Tuple, TypeAlias
 from renderer.types import FaceIndices, Normals, Texture, UVCoordinates, Vertices
 
-jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
+if hasattr(jax.config, "jax_array"):
+    jax.config.update("jax_array", True)  # pyright: ignore[reportUnknownMemberType]
 
 T2f: TypeAlias = Tuple[float, float]
 T3f: TypeAlias = Tuple[float, float, float]