How to use a develop branch of IREE with iree-jax

[Yi Wang]

I built IREE from the source code in this feature branch https://github.com/antiagainst/iree/tree/apple-metal-hal. Then, I want to install and run models/gpt2.

I following the step here https://github.com/antiagainst/iree/tree/apple-metal-hal, and run the dependency installation commands in a fresh new Conda enviromnet with Python 3.9. I didn't use --use-pinned because I am using a feature branch. And I have PYTHONPATH set by ~/w/iree/build/.env.

cd ~/w/iree-jax
python3 -m pip install absl-py future h5py numpy optax transformers # This works
python build_tools/configure.py # This updates version_info.json
python build_tools/pip_install.py # This failed

The last command failed with the following error message

10:59 $ python3 build_tools/pip_install.py
20230331.475
python3 -m pip install . -f https://openxla.github.io/iree/pip-release-links.html --force-reinstall .[xla,cpu,test]
Looking in indexes: https://pypi.apple.com/simple
Looking in links: https://openxla.github.io/iree/pip-release-links.html
Processing /Users/y/w/iree-jax
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done
Collecting numpy (from iree-jax==0.1.dev1)
  Using cached https://pypi.apple.com/packages/packages/54/41/fb17c1d48a574c50422ff3f1b17ed979b755adc6ed291c4a44a76e226c67/numpy-1.24.3-cp39-cp39-macosx_11_0_arm64.whl (13.9 MB)
Collecting jax>=0.4.8 (from iree-jax==0.1.dev1)
  Using cached jax-0.4.8-py3-none-any.whl
INFO: pip is looking at multiple versions of iree-jax to determine which version is compatible with other requirements. This could take a while.
ERROR: Could not find a version that satisfies the requirement iree-compiler>=20230331.475 (from iree-jax) (from versions: 20220420.113, 20220421.114, 20220422.115, 20220423.116, 20220424.117, 20220425.118, 20220426.119, 20220427.120, 20220428.121, 20220429.122, 20220430.123, 20220501.124, 20220502.125, 20220503.126, 20220504.127, 20220505.128, 20220506.129, 20220507.130, 20220508.131, 20220509.132, 20220513.137, 20220514.138, 20220515.139, 20220516.140, 20220517.141, 20220518.142, 20220519.143, 20220520.144, 20220522.146, 20220523.147, 20220525.149, 20220526.150, 20220527.151, 20220601.156, 20220602.157, 20220603.158, 20220604.159, 20220605.160, 20220606.161, 20220607.162, 20220608.163, 20220609.164, 20220611.166, 20220612.167, 20220614.169, 20220615.172, 20220616.173, 20220616.174, 20220617.175, 20220618.176, 20220619.177, 20220620.178, 20220621.179, 20220622.180, 20220623.181, 20220624.182, 20220625.183, 20220626.184, 20220627.185, 20220629.188, 20220630.189, 20220701.191, 20220702.192, 20220703.193, 20220704.194, 20220705.195, 20220706.196, 20220707.197, 20220708.198, 20220709.199, 20220710.200, 20220712.202, 20220713.203, 20220714.204, 20220715.205, 20220716.206, 20220717.207, 20220718.208, 20220720.210, 20220721.211, 20220730.220, 20220731.221, 20220801.222, 20220803.224, 20220809.230, 20220810.231, 20220811.232, 20220812.233, 20220813.234, 20220814.235, 20220816.237, 20220817.238, 20220819.240, 20220820.241, 20220821.242, 20220823.244, 20220825.246, 20220826.247, 20220828.249, 20220830.251, 20220831.252, 20220901.253, 20220902.254, 20220903.255, 20220904.256, 20220905.257, 20220906.258, 20220907.259, 20220908.260, 20220909.261, 20220910.262, 20220911.263, 20220912.264, 20220913.265, 20220914.266, 20220915.267, 20220917.269, 20220918.270, 20220919.5, 20220919.6, 20220919.271, 20220920.272, 20220921.273, 20220922.274, 20220923.275, 20220924.276, 20220925.277, 20220926.278, 20220928.280, 20220929.7, 20220929.281, 20220930.282, 20221001.283, 20221002.284, 20221003.285, 20221004.286, 20221006.288, 20221007.289, 20221008.290, 20221009.291, 20221010.292, 20221011.293, 20221012.294, 20221019.301, 20221020.302, 20221021.303, 20221022.304, 20221023.305, 20221024.306, 20221025.307, 20221026.308, 20221027.309, 20221028.310, 20221029.311, 20221030.312, 20221031.313, 20221101.314, 20221102.315, 20221103.316, 20221104.317, 20221105.318, 20221106.319, 20221107.320, 20221108.321, 20221109.322, 20221110.323, 20221111.324, 20221112.325, 20221113.326, 20221114.327, 20221115.328, 20221116.329, 20221117.330, 20221118.331, 20221119.332, 20221120.333, 20221121.334, 20221122.335, 20221123.336, 20221124.337, 20221126.10, 20221127.340, 20221129.342, 20221201.344, 20221202.345, 20221203.346, 20221204.347, 20221205.348, 20221207.350, 20221208.11, 20221208.351, 20221209.353, 20221210.354, 20221211.355, 20221212.356, 20221213.357, 20221214.358, 20221215.359, 20221216.360, 20221217.361, 20221218.362, 20221219.363, 20221223.368, 20221224.369, 20221225.370, 20221226.371, 20221227.372, 20221231.376, 20230101.377, 20230102.378, 20230104.386, 20230105.387, 20230106.389, 20230107.390, 20230108.391, 20230110.393, 20230111.394, 20230112.395, 20230114.397, 20230115.398, 20230117.400, 20230119.402, 20230120.403, 20230121.404, 20230122.405, 20230123.406, 20230125.408, 20230126.409, 20230127.410, 20230128.411, 20230129.412, 20230130.413, 20230131.414, 20230201.415, 20230202.416, 20230203.417, 20230204.418, 20230205.420, 20230205.421, 20230206.422, 20230207.423, 20230208.424, 20230209.425, 20230210.426, 20230211.427, 20230212.428)
ERROR: No matching distribution found for iree-compiler>=20230331.475
None

Despite the above errors, python3 -c import iree.jax works. However, running the script models/gpt2/export.py complains no module named iree.jax.

(iree-jax) ✔ ~/w/iree-jax [main|✚ 2]
11:12 $ python3 models/gpt2/export.py
Traceback (most recent call last):
  File "/Users/y/w/iree-jax/models/gpt2/export.py", line 16, in <module>
    import iree.jax
ModuleNotFoundError: No module named 'iree.jax'

[Yi Wang]

Sorry for this false alarm - I did't set PYTHONPATH correctly to point to iree.runtime, which I built from source code.