[BUG] optree on Windows on ARM64

Required prerequisites

• I have read the documentation https://optree.readthedocs.io.
• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

What version of OpTree are you using?

0.13.1

System information

Installation: pip (please check problem description)

Python: 3.12.8 (Windows ARM64)

Platform: Windows ARM64

Optree Version: 0.13.1

Problem description

Windows ARM64 version of optree has a problem on import.

optree-0.13.1-cp313-cp313-win_arm64.whl

If you install it with regular command (by wheel file):

pip install optree

or to prevent any cache

pip install optree --no-cache-dir

ImportError: cannot import name '_C' from partially initialized module 'optree' (most likely due to a circular import)

---

If you install it via source code (Visual Studio C++ Build Tools are installed in system):

pip install --no-binary optree optree

or

pip install --no-binary optree optree --no-cache-dir

It is successfully being installed and working without any problem.

Reproducible example code

import optree

Traceback

Traceback (most recent call last):
  File "main.py", line 1, in <module>
    import optree
  File "__init__.py", line 17, in <module>
    from optree import accessor, dataclasses, functools, integration, typing
  File "site-packages\optree\accessor.py", line 36, in <module>
    from optree import _C
ImportError: cannot import name '_C' from partially initialized module 'optree' (most likely due to a circular import) (optree\__init__.py)

Expected behavior

It should import optree without any problem.

Additional context

Most probably there is a problem on generation of Windows ARM64 wheel.

Until fixing the issue, as a workaround, you might also consider to remove optree-0.13.1-cp313-cp313-win_arm64.whl from wheel list to force build on Windows ARM64.

[XuehaiPan]

Hi @ozanMSFT, thanks for raising this. But I cannot reproduce your error in a fresh venv environment.

#  % PanXuehai @ Windows-ARM64 in ~ [20:26:05]
$ python3.13 -m venv venv
#  % PanXuehai @ Windows-ARM64 in ~ [20:26:21]
$ & .\venv\Scripts\Activate.ps1
(optree )
#  % PanXuehai @ Windows-ARM64 in ~ [20:26:26]
$ pip3 install -U pip setuptools wheel build ipython
...
(optree )
#  % PanXuehai @ Windows-ARM64 in ~\Projects [20:27:56]
$ pip3 install optree
Collecting optree
  Downloading optree-0.13.1-cp313-cp313-win_amd64.whl.metadata (48 kB)
Requirement already satisfied: typing-extensions>=4.5.0 in c:\users\panxuehai\projects\optree\venv\lib\site-packages (from optree) (4.12.2)
Downloading optree-0.13.1-cp313-cp313-win_amd64.whl (293 kB)
Installing collected packages: optree
Successfully installed optree-0.13.1
(optree )
#  % PanXuehai @ Windows-ARM64 in ~ [20:28:07]
$ ipython

Python 3.13.1 (tags/v3.13.1:0671451, Dec  3 2024, 19:06:28) [MSC v.1942 64 bit (AMD64)]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.31.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import optree

In [2]: optree.__version__
Out[2]: '0.13.1'

In [3]: import platform

In [4]: platform.system()
Out[4]: 'Windows'

In [5]: platform.uname()
Out[5]: uname_result(system='Windows', node='Windows-ARM64', release='11', version='10.0.26100', machine='ARM64')

In [6]: platform.python_version()
Out[6]: '3.13.1'

In [7]: platform.python_implementation()
Out[7]: 'CPython'

In [8]: import sys

In [9]: sys.implementation
Out[9]:
namespace(name='cpython',
          cache_tag='cpython-313',
          version=sys.version_info(major=3, minor=13, micro=1, releaselevel='final', serial=0),
          hexversion=51184112)

[XuehaiPan]

Downloading optree-0.13.1-cp313-cp313-win_amd64.whl (293 kB)

I see pip loads the wheel with the tag win_amd64. The OS is an ARM64 Windows virtual machine hosted in macOS with Apple Silicon chip. Did I miss something?

Thanks for the reply @XuehaiPan

Based on the logs you've provided, I'm seeing that it is downloading x64 version in your system:

Downloading optree-0.13.1-cp313-cp313-win_amd64.whl.metadata (48 kB)

However, it should download optree-0.13.1-cp313-cp313-win_arm64.whl

---

It might be because of some caching in your environment.

Could you please clear your pip cache or use following command:

pip install optree --no-cache-dir

It should print:

Downloading optree-0.13.1-cp312-cp312-win_arm64.whl.metadata (48 kB)

---

I'm using Python 3.12.8

[XuehaiPan]

I have tried a new fresh environment with Python 3.12. It still installs the wheel with the tag win_amd64.

#  % PanXuehai @ Windows-ARM64 in ~ [20:46:14]
$ mkdir test

    Directory: C:\Users\PanXuehai

Mode                 LastWriteTime         Length Name
----                 -------------         ------ ----
d----           2025/1/14    20:46                test

#  % PanXuehai @ Windows-ARM64 in ~ [20:46:16]
$ cd test
#  % PanXuehai @ Windows-ARM64 in ~\test [20:46:17]
$ python3.12 -m venv venv
#  % PanXuehai @ Windows-ARM64 in ~\test [20:46:31]
$ & .\venv\Scripts\Activate.ps1
(test )
#  % PanXuehai @ Windows-ARM64 in ~\test [20:46:37]
$ pip3 install -U pip setuptools wheel build ipython
...
(test )
#  % PanXuehai @ Windows-ARM64 in ~\test [20:47:28]
$ pip3 cache purge
Files removed: 116
(test )
#  % PanXuehai @ Windows-ARM64 in ~\test [20:47:35]
$ pip3 install optree --no-cache-dir
Collecting optree
  Downloading optree-0.13.1-cp312-cp312-win_amd64.whl.metadata (48 kB)
Collecting typing-extensions>=4.5.0 (from optree)
  Downloading typing_extensions-4.12.2-py3-none-any.whl.metadata (3.0 kB)
Downloading optree-0.13.1-cp312-cp312-win_amd64.whl (292 kB)
Downloading typing_extensions-4.12.2-py3-none-any.whl (37 kB)
Installing collected packages: typing-extensions, optree
Successfully installed optree-0.13.1 typing-extensions-4.12.2
(test )
#  % PanXuehai @ Windows-ARM64 in ~\test [20:47:46]
$ ipython
Python 3.12.8 (tags/v3.12.8:2dc476b, Dec  3 2024, 19:30:04) [MSC v.1942 64 bit (AMD64)]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.31.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import optree

In [2]: optree.__version__
Out[2]: '0.13.1'

In [3]: import platform

In [4]: platform.uname()
Out[4]: uname_result(system='Windows', node='Windows-ARM64', release='11', version='10.0.26100', machine='ARM64')

In [5]: import sys

In [6]: sys.implementation
Out[6]:
namespace(name='cpython',
          cache_tag='cpython-312',
          version=sys.version_info(major=3, minor=12, micro=8, releaselevel='final', serial=0),
          hexversion=51120368)

I think problem might be also related to your Python version.

this is your output:

Python 3.12.8 (tags/v3.12.8:2dc476b, Dec 3 2024, 19:30:04) [MSC v.1942 64 bit (AMD64)]

this is Windows ARM64 output:

Python 3.12.8 (tags/v3.12.8:2dc476b, Dec 3 2024, 19:45:50) [MSC v.1942 64 bit (ARM64)]

---

Please try to install:

https://www.python.org/downloads/windows/

Windows installer (ARM64)

[XuehaiPan]

If you download the ARM64 wheel file and unzip it, the C extension is tagged with win_amd64. That's why your Python interpreter has not found the _C extension. It might be a bug for cibuildwheel for Windows cross-compiling.

(test )
#  % PanXuehai @ Windows-ARM64 in ~\test [20:51:10] C:1
$ wget https://files.pythonhosted.org/packages/80/34/d1b1849a6240385c4a3af5da9425b11912204d0b1cf142d802815319b73a/optree-0.13.1-cp313-cp313-win_arm64.whl
--2025-01-14 20:53:44--  https://files.pythonhosted.org/packages/80/34/d1b1849a6240385c4a3af5da9425b11912204d0b1cf142d802815319b73a/optree-0.13.1-cp313-cp313-win_arm64.whl
Connecting to 127.0.0.1:7890... connected.
Proxy request sent, awaiting response... 200 OK
Length: 293670 (287K) [application/octet-stream]
Saving to: 'optree-0.13.1-cp313-cp313-win_arm64.whl'

optree-0.13.1-cp313-cp313-win_arm64.whl              100%[===================================================================================================================>] 286.79K   819KB/s    in 0.4s

2025-01-14 20:53:44 (819 KB/s) - 'optree-0.13.1-cp313-cp313-win_arm64.whl' saved [293670/293670]

(test )
#  % PanXuehai @ Windows-ARM64 in ~\test [20:56:18]
$ wheel unpack --dest .\sitedir .\optree-0.13.1-cp313-cp313-win_arm64.whl
Unpacking to: sitedir\optree-0.13.1...OK
(test )
#  % PanXuehai @ Windows-ARM64 in ~\test [20:56:52]
$ Get-ChildItem .\sitedir\optree-0.13.1\optree\

    Directory: C:\Users\PanXuehai\test\sitedir\optree-0.13.1\optree

Mode                 LastWriteTime         Length Name
----                 -------------         ------ ----
d----           2025/1/14    20:56                integration
-a---           2025/1/14    20:56           5821 __init__.py
-a---           2025/1/14    20:56         550400 _C.cp313-win_amd64.pyd
-a---           2025/1/14    20:56           5520 _C.pyi
-a---           2025/1/14    20:56          14851 accessor.py
-a---           2025/1/14    20:56          18313 dataclasses.py
-a---           2025/1/14    20:56           6831 functools.py
-a---           2025/1/14    20:56         149248 ops.py
-a---           2025/1/14    20:56              0 py.typed
-a---           2025/1/14    20:56          31663 registry.py
-a---           2025/1/14    20:56          18684 typing.py
-a---           2025/1/14    20:56           3455 utils.py
-a---           2025/1/14    20:56           1880 version.py

[XuehaiPan]

Maybe related:

• [Windows ARM] sys_tags produces win_amd64 pypa/cibuildwheel#1942
• [BUG] win-arm64 platform is not detected pypa/setuptools#2724
• Fixes #2724 Adds detection of Windows ARM64 build environments pypa/setuptools#2725
• Fix setup.py to detect cross-compiling to ARM64 on Windows python-cffi/cffi#101

I installed the Python ARM64 build. I renamed the file extension in site-packages and I got:

In [1]: import _C
---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[1], line 1
----> 1 import _C

ImportError: DLL load failed while importing _C: %1 is not valid Win32 application

[XuehaiPan]

Hi @ozanMSFT, I opened a PR to fix this.

• fix(setup): fix cross-compiling for ARM64 on x64 Windows #183

The built wheels for ARM64 Windows work fine:

# PanXuehai @ Windows-ARM64 in ~\Downloads [20:14:24]
$ Expand-Archive -Force .\wheels-cp313-Windows-ARM64.zip wheels
# PanXuehai @ Windows-ARM64 in ~\Downloads [20:14:28]
$ cd .\wheels\
# PanXuehai @ Windows-ARM64 in ~\Downloads\wheels [20:14:36]
$ python3.13 -m venv venv
# PanXuehai @ Windows-ARM64 in ~\Downloads\wheels [20:14:49]
$ & .\venv\Scripts\Activate.ps1
(wheels )
# PanXuehai @ Windows-ARM64 in ~\Downloads\wheels [20:14:55]
$ pip3 install .\optree-0.13.1-cp313-cp313-win_arm64.whl
Processing c:\users\panxuehai\downloads\wheels\optree-0.13.1-cp313-cp313-win_arm64.whl
Collecting typing-extensions>=4.5.0 (from optree==0.13.1)
  Using cached typing_extensions-4.12.2-py3-none-any.whl.metadata (3.0 kB)
Using cached typing_extensions-4.12.2-py3-none-any.whl (37 kB)
Installing collected packages: typing-extensions, optree
Successfully installed optree-0.13.1 typing-extensions-4.12.2
(wheels )
# PanXuehai @ Windows-ARM64 in ~\Downloads\wheels [20:15:13]
$ python
Python 3.13.1 (tags/v3.13.1:0671451, Dec  3 2024, 20:04:37) [MSC v.1942 64 bit (ARM64)] on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import optree
>>> optree.__version__
'0.13.1'
>>> optree.tree_map(lambda x: x**2, {'a': [1, 2], 'b': (3, {'c': 4}, 5)})
{'a': [1, 4], 'b': (9, {'c': 16}, 25)}

You can view the wheels at https://github.com/metaopt/optree/actions/runs/12806409239?pr=183

I can confirm that it is fixed in version 0.14.0.

@XuehaiPan, thanks a lot for the quick fix!