PyTreeDef comparison fails type checking for jax==0.6.1

[mfschubert]

Description

Starting with jax 0.6.1, it seems that the ability to compare PyTreeDef is not recognized by mypy. Specifically, the following file:

from typing import Any
from jax import tree_util

def test_pytreedef_equal(a: Any, b: Any) -> bool:
  return tree_util.tree_structure(a) == tree_util.tree_structure(b)

will fail mypy type validation with the following error:

test.py:6: error: Unsupported left operand type for == (PyTreeDef?)  [operator]
Found 1 error in 1 file (checked 1 source file)

This is demonstrated in the following colab.

The error does not occur for jax 0.6.0 and earlier.

System info (python version, jaxlib version, accelerator, etc.)

jax: 0.6.1
jaxlib: 0.6.1
numpy: 2.0.2
python: 3.11.12 (main, Apr 9 2025, 08:55:54) [GCC 11.4.0]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='44b407801e70', release='6.1.123+', version='#1 SMP PREEMPT_DYNAMIC Sun Mar 30 16:01:29 UTC 2025', machine='x86_64')

[jakevdp]

Somehow it seems mypy is not discovering the type declaration here:

jax/jaxlib/_jax/pytree.pyi

Line 134 in 6f0b993

def __eq__(self, __other: Any) -> bool: ...

The difference in v0.6.1 is that this module was moved from jaxlib.xla_extension.pytree to jaxlib._jax.pytree in #28181 I'm not sure why that would lead to an error here...

[jakevdp]

I'm going to assign @superbobry who may have some insight

[hawkinsp]

We did stop shipping the type stubs as part of jaxlib, that is probably why.

(The thinking here was: jaxlib is private, and it's easier to just keep the stubs as part of jax.)