Fix for jax 0.8.2 (#455)

* fix for jax 0.8.2

* add more register

* format

Author: jfeser

docs/source/beam.py

@@ -6,7 +6,6 @@
 import functools
 import heapq
 import random
-import typing
 from collections.abc import Callable
 from dataclasses import dataclass
 from pprint import pprint

effectful/handlers/jax/_terms.py

@@ -432,7 +432,8 @@ def ndim(self) -> int:
         return len(self.shape)
 
 
-@bind_dims.register  # type: ignore
+@bind_dims.register(jax.Array)  # type: ignore
+@bind_dims.register(jax._src.core.Tracer)  # type: ignore
 def _bind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array:
     """Convert named dimensions to positional dimensions.
 
@@ -501,6 +502,7 @@ def _evaluate(expr):
     return reindexed
 
 
-@unbind_dims.register  # type: ignore
+@unbind_dims.register(jax.Array)  # type: ignore
+@unbind_dims.register(jax._src.core.Tracer)  # type: ignore
 def _unbind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array:
     return jax_getitem(t, tuple(n() for n in args))