Fix labels, pyright in CI (#1592)

Author: sritchie

.github/workflows/labeler.yml

@@ -8,6 +8,9 @@ on:
 jobs:
   labeler:
     runs-on: ubuntu-22.04
+    permissions:
+      contents: read
+      issues: write
     steps:
       - name: Check out the repository
         uses: actions/checkout@v4

.github/workflows/ruff.yml

@@ -18,5 +18,5 @@ jobs:
 
     - uses: chartboost/ruff-action@v1
       with:
-        version: 0.9.9
+        version: 0.11.2
         args: check --output-format github

.pre-commit-config.yaml

@@ -43,6 +43,6 @@ repos:
       - id: vulture
 
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.398
+    rev: v1.1.399
     hooks:
     - id: pyright

src/genjax/_src/adev/primitives.py

@@ -11,8 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 """Defines ADEV primitives."""
 
+# pyright: reportPrivateImportUsage=false
+
 import jax
 import jax._src.core
 import jax._src.dtypes as jax_dtypes

src/genjax/_src/core/compiler/staging.py

@@ -287,8 +287,8 @@ def stage(f):
     """Returns a function that stages a function to a ClosedJaxpr."""
 
     def wrapped(*args, **kwargs):
-        debug_info = api_util.debug_info("Tracing to Jaxpr", f, args, kwargs)
-        fun = lu.wrap_init(f, params=kwargs, debug_info=debug_info)
+        debug_info = api_util.debug_info("Tracing to Jaxpr", f, args, kwargs)  # pyright: ignore[reportAttributeAccessIssue]
+        fun = lu.wrap_init(f, params=kwargs, debug_info=debug_info)  # pyright: ignore[reportCallIssue]
         flat_args, in_tree = jtu.tree_flatten(args)
         flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
         flat_avals = safe_map(get_shaped_aval, flat_args)

src/genjax/_src/core/generative/choice_map.py

@@ -1493,8 +1493,7 @@ def build(chm: ChoiceMap, addr: DynamicAddressComponent) -> ChoiceMap:
             return Indexed(chm, addr)
 
     def filter(self, selection: Selection | Flag) -> ChoiceMap:
-        addr = _full_slice if self.addr is None else self.addr
-        return self.c.filter(selection).extend(addr)
+        return self.c.filter(selection).extend(self.addr)
 
     def get_value(self) -> Any:
         return None
@@ -1510,13 +1509,7 @@ def get_inner_map(self, addr: AddressComponent) -> ChoiceMap:
                     "Only scalar dynamic addresses are supported by get_submap."
                 )
 
-            if self.addr is None:
-                # None means that this instance was created with `:`, so no masking is required and we assume that the user will provide an in-bounds `int | ScalarInt`` address. If they don't they will run up against JAX's clamping behavior.
-                return jtu.tree_map(
-                    lambda v: v[addr], self.c, is_leaf=lambda x: isinstance(x, Mask)
-                )
-
-            elif isinstance(self.addr, Array) and self.addr.shape:
+            if isinstance(self.addr, Array) and self.addr.shape:
                 # We can't allow slices, as self.addr might look like, e.g. `[2,5,6]`, and we don't have any way to combine this "sparse array selector" with an incoming slice.
                 assert not isinstance(addr, slice), (
                     f"Slices are not allowed against array-shaped dynamic addresses. Tried to apply {addr} to {self.addr}."