cyxla 0.1.0__tar.gz

cyxla-0.1.0/.gitignore

@@ -0,0 +1,18 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv*
+
+# Test output
+.coverage
+.coverage.*
+coverage.xml
+
+# Coverage builds copy the generated C file next to the Cython source.
+/src/cyxla/*.c

cyxla-0.1.0/.python-version

@@ -0,0 +1 @@
+3.11

cyxla-0.1.0/LICENSE.md

@@ -0,0 +1,135 @@
+**HIPPOCRATIC LICENSE**
+
+**Version 3.0, October 2021**
+
+<https://firstdonoharm.dev/version/3/0/bds-eco-ffd-law-mil-sv.md>
+
+**TERMS AND CONDITIONS**
+
+TERMS AND CONDITIONS FOR USE, COPY, MODIFICATION, PREPARATION OF DERIVATIVE WORK, REPRODUCTION, AND DISTRIBUTION:
+
+**[1.](#1) DEFINITIONS:**
+
+_This section defines certain terms used throughout this license agreement._
+
+[1.1.](#1.1) “License” means the terms and conditions, as stated herein, for use, copy, modification, preparation of derivative work, reproduction, and distribution of Software (as defined below).
+
+[1.2.](#1.2) “Licensor” means the copyright and/or patent owner or entity authorized by the copyright and/or patent owner that is granting the License.
+
+[1.3.](#1.3) “Licensee” means the individual or entity exercising permissions granted by this License, including the use, copy, modification, preparation of derivative work, reproduction, and distribution of Software (as defined below).
+
+[1.4.](#1.4) “Software” means any copyrighted work, including but not limited to software code, authored by Licensor and made available under this License.
+
+[1.5.](#1.5) “Supply Chain” means the sequence of processes involved in the production and/or distribution of a commodity, good, or service offered by the Licensee.
+
+[1.6.](#1.6) “Supply Chain Impacted Party” or “Supply Chain Impacted Parties” means any person(s) directly impacted by any of Licensee’s Supply Chain, including the practices of all persons or entities within the Supply Chain prior to a good or service reaching the Licensee.
+
+[1.7.](#1.7) “Duty of Care” is defined by its use in tort law, delict law, and/or similar bodies of law closely related to tort and/or delict law, including without limitation, a requirement to act with the watchfulness, attention, caution, and prudence that a reasonable person in the same or similar circumstances would use towards any Supply Chain Impacted Party.
+
+[1.8.](#1.8) “Worker” is defined to include any and all permanent, temporary, and agency workers, as well as piece-rate, salaried, hourly paid, legal young (minors), part-time, night, and migrant workers.
+
+**[2.](#2) INTELLECTUAL PROPERTY GRANTS:**
+
+_This section identifies intellectual property rights granted to a Licensee_.
+
+[2.1.](#2.1) _Grant of Copyright License_: Subject to the terms and conditions of this License, Licensor hereby grants to Licensee a worldwide, non-exclusive, no-charge, royalty-free copyright license to use, copy, modify, prepare derivative work, reproduce, or distribute the Software, Licensor authored modified software, or other work derived from the Software.
+
+[2.2.](#2.2) _Grant of Patent License_: Subject to the terms and conditions of this License, Licensor hereby grants Licensee a worldwide, non-exclusive, no-charge, royalty-free patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer Software.
+
+**[3.](#3) ETHICAL STANDARDS:**
+
+_This section lists conditions the Licensee must comply with in order to have rights under this License._
+
+The rights granted to the Licensee by this License are expressly made subject to the Licensee’s ongoing compliance with the following conditions:
+
+* [3.1.](#3.1) The Licensee SHALL NOT, whether directly or indirectly, through agents or assigns:
+
+  * [3.1.1.](#3.1.1) Infringe upon any person’s right to life or security of person, engage in extrajudicial killings, or commit murder, without lawful cause (See Article 3, _United Nations Universal Declaration of Human Rights_; Article 6, _International Covenant on Civil and Political Rights_)
+  * [3.1.2.](#3.1.2) Hold any person in slavery, servitude, or forced labor (See Article 4, _United Nations Universal Declaration of Human Rights_; Article 8, _International Covenant on Civil and Political Rights_);
+  * [3.1.3.](#3.1.3) Contribute to the institution of slavery, slave trading, forced labor, or unlawful child labor (See Article 4, _United Nations Universal Declaration of Human Rights_; Article 8, _International Covenant on Civil and Political Rights_);
+  * [3.1.4.](#3.1.4) Torture or subject any person to cruel, inhumane, or degrading treatment or punishment (See Article 5, _United Nations Universal Declaration of Human Rights_; Article 7, _International Covenant on Civil and Political Rights_);
+  * [3.1.5.](#3.1.5) Discriminate on the basis of sex, gender, sexual orientation, race, ethnicity, nationality, religion, caste, age, medical disability or impairment, and/or any other like circumstances (See Article 7, _United Nations Universal Declaration of Human Rights_; Article 2, _International Covenant on Economic, Social and Cultural Rights_; Article 26, _International Covenant on Civil and Political Rights_);
+  * [3.1.6.](#3.1.6) Prevent any person from exercising his/her/their right to seek an effective remedy by a competent court or national tribunal (including domestic judicial systems, international courts, arbitration bodies, and other adjudicating bodies) for actions violating the fundamental rights granted to him/her/them by applicable constitutions, applicable laws, or by this License (See Article 8, _United Nations Universal Declaration of Human Rights_; Articles 9 and 14, _International Covenant on Civil and Political Rights_);
+  * [3.1.7.](#3.1.7) Subject any person to arbitrary arrest, detention, or exile (See Article 9, _United Nations Universal Declaration of Human Rights_; Article 9, _International Covenant on Civil and Political Rights_);
+  * [3.1.8.](#3.1.8) Subject any person to arbitrary interference with a person’s privacy, family, home, or correspondence without the express written consent of the person (See Article 12, _United Nations Universal Declaration of Human Rights_; Article 17, _International Covenant on Civil and Political Rights_);
+  * [3.1.9.](#3.1.9) Arbitrarily deprive any person of his/her/their property (See Article 17, _United Nations Universal Declaration of Human Rights_);
+  * [3.1.10.](#3.1.10) Forcibly remove indigenous peoples from their lands or territories or take any action with the aim or effect of dispossessing indigenous peoples from their lands, territories, or resources, including without limitation the intellectual property or traditional knowledge of indigenous peoples, without the free, prior, and informed consent of indigenous peoples concerned (See Articles 8 and 10, _United Nations Declaration on the Rights of Indigenous Peoples_);
+  * [3.1.11.](#3.1.11) _Fossil Fuel Divestment_: Be an individual or entity, or a representative, agent, affiliate, successor, attorney, or assign of an individual or entity, on the [FFI Solutions Carbon Underground 200 list](https://www.ffisolutions.com/research-analytics-index-solutions/research-screening/the-carbon-underground-200/?cn-reloaded=1);
+  * [3.1.12.](#3.1.12) _Ecocide_: Commit ecocide:
+
+    * [3.1.12.1.](#3.1.12.1) For the purpose of this section, “ecocide” means unlawful or wanton acts committed with knowledge that there is a substantial likelihood of severe and either widespread or long-term damage to the environment being caused by those acts;
+    * [3.1.12.2.](#3.1.12.2) For the purpose of further defining ecocide and the terms contained in the previous paragraph:
+
+      * [3.1.12.2.1.](#3.1.12.2.1) “Wanton” means with reckless disregard for damage which would be clearly excessive in relation to the social and economic benefits anticipated;
+      * [3.1.12.2.2.](#3.1.12.2.2) “Severe” means damage which involves very serious adverse changes, disruption, or harm to any element of the environment, including grave impacts on human life or natural, cultural, or economic resources;
+      * [3.1.12.2.3.](#3.1.12.2.3) “Widespread” means damage which extends beyond a limited geographic area, crosses state boundaries, or is suffered by an entire ecosystem or species or a large number of human beings;
+      * [3.1.12.2.4.](#3.1.12.2.4) “Long-term” means damage which is irreversible or which cannot be redressed through natural recovery within a reasonable period of time; and
+      * [3.1.12.2.5.](#3.1.12.2.5) “Environment” means the earth, its biosphere, cryosphere, lithosphere, hydrosphere, and atmosphere, as well as outer space  
+            (See Section II, _Independent Expert Panel for the Legal Definition of Ecocide_, Stop Ecocide Foundation and the Promise Institute for Human Rights at UCLA School of Law, June 2021);
+  * [3.1.13.](#3.1.13) _Boycott / Divestment / Sanctions_: Be an individual or entity, or a representative, agent, affiliate, successor, attorney, or assign of an individual or entity, identified by the Boycott, Divestment, Sanctions (“BDS”) movement on its website (<https://bdsmovement.net/> and <https://bdsmovement.net/get-involved/what-to-boycott>) as a target for boycott;
+  * [3.1.14.](#3.1.14) _Mass Surveillance_: Be a government agency or multinational corporation, or a representative, agent, affiliate, successor, attorney, or assign of a government or multinational corporation, which participates in mass surveillance programs;
+  * [3.1.15.](#3.1.15) _Military Activities_: Be an entity or a representative, agent, affiliate, successor, attorney, or assign of an entity which conducts military activities;
+  * [3.1.16.](#3.1.16) _Law Enforcement_: Be an individual or entity, or a representative, agent, affiliate, successor, attorney, or assign of an individual or entity, that provides good or services to, or otherwise enters into any commercial contracts with, any local, state, or federal law enforcement agency;
+  * [3.1.17.](#3.1.17) Interfere with Workers’ free exercise of the right to organize and associate (See Article 20, United Nations Universal Declaration of Human Rights; C087 - Freedom of Association and Protection of the Right to Organise Convention, 1948 (No. 87), International Labour Organization; Article 8, International Covenant on Economic, Social and Cultural Rights); and
+  * [3.1.18.](#3.1.18) Harm the environment in a manner inconsistent with local, state, national, or international law.
+* [3.2.](#3.2) The Licensee SHALL:
+
+  * [3.2.1.](#3.2.1) Provide equal pay for equal work where the performance of such work requires equal skill, effort, and responsibility, and which are performed under similar working conditions, except where such payment is made pursuant to:
+
+    * [3.2.1.1.](#3.2.1.1) A seniority system;
+    * [3.2.1.2.](#3.2.1.2) A merit system;
+    * [3.2.1.3.](#3.2.1.3) A system which measures earnings by quantity or quality of production; or
+    * [3.2.1.4.](#3.2.1.4) A differential based on any other factor other than sex, gender, sexual orientation, race, ethnicity, nationality, religion, caste, age, medical disability or impairment, and/or any other like circumstances (See 29 U.S.C.A. § 206(d)(1); Article 23, _United Nations Universal Declaration of Human Rights_; Article 7, _International Covenant on Economic, Social and Cultural Rights_; Article 26, _International Covenant on Civil and Political Rights_); and
+  * [3.2.2.](#3.2.2) Allow for reasonable limitation of working hours and periodic holidays with pay (See Article 24, _United Nations Universal Declaration of Human Rights_; Article 7, _International Covenant on Economic, Social and Cultural Rights_).
+
+**[4.](#4) SUPPLY CHAIN IMPACTED PARTIES:**
+
+_This section identifies additional individuals or entities that a Licensee could harm as a result of violating the Ethical Standards section, the condition that the Licensee must voluntarily accept a Duty of Care for those individuals or entities, and the right to a private right of action that those individuals or entities possess as a result of violations of the Ethical Standards section._
+
+[4.1.](#4.1) In addition to the above Ethical Standards, Licensee voluntarily accepts a Duty of Care for Supply Chain Impacted Parties of this License, including individuals and communities impacted by violations of the Ethical Standards. The Duty of Care is breached when a provision within the Ethical Standards section is violated by a Licensee, one of its successors or assigns, or by an individual or entity that exists within the Supply Chain prior to a good or service reaching the Licensee.
+
+[4.2.](#4.2) Breaches of the Duty of Care, as stated within this section, shall create a private right of action, allowing any Supply Chain Impacted Party harmed by the Licensee to take legal action against the Licensee in accordance with applicable negligence laws, whether they be in tort law, delict law, and/or similar bodies of law closely related to tort and/or delict law, regardless if Licensee is directly responsible for the harms suffered by a Supply Chain Impacted Party. Nothing in this section shall be interpreted to include acts committed by individuals outside of the scope of his/her/their employment.
+
+[5.](#5) **NOTICE:** _This section explains when a Licensee must notify others of the License._
+
+[5.1.](#5.1) _Distribution of Notice_: Licensee must ensure that everyone who receives a copy of or uses any part of Software from Licensee, with or without changes, also receives the License and the copyright notice included with Software (and if included by the Licensor, patent, trademark, and attribution notice). Licensee must ensure that License is prominently displayed so that any individual or entity seeking to download, copy, use, or otherwise receive any part of Software from Licensee is notified of this License and its terms and conditions. Licensee must cause any modified versions of the Software to carry prominent notices stating that Licensee changed the Software.
+
+[5.2.](#5.2) _Modified Software_: Licensee is free to create modifications of the Software and distribute only the modified portion created by Licensee, however, any derivative work stemming from the Software or its code must be distributed pursuant to this License, including this Notice provision.
+
+[5.3.](#5.3) _Recipients as Licensees_: Any individual or entity that uses, copies, modifies, reproduces, distributes, or prepares derivative work based upon the Software, all or part of the Software’s code, or a derivative work developed by using the Software, including a portion of its code, is a Licensee as defined above and is subject to the terms and conditions of this License.
+
+**[6.](#6) REPRESENTATIONS AND WARRANTIES:**
+
+[6.1.](#6.1) _Disclaimer of Warranty_: TO THE FULL EXTENT ALLOWED BY LAW, THIS SOFTWARE COMES “AS IS,” WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED, AND LICENSOR SHALL NOT BE LIABLE TO ANY PERSON OR ENTITY FOR ANY DAMAGES OR OTHER LIABILITY ARISING FROM, OUT OF, OR IN CONNECTION WITH THE SOFTWARE OR THIS LICENSE, UNDER ANY LEGAL CLAIM.
+
+[6.2.](#6.2) _Limitation of Liability_: LICENSEE SHALL HOLD LICENSOR HARMLESS AGAINST ANY AND ALL CLAIMS, DEBTS, DUES, LIABILITIES, LIENS, CAUSES OF ACTION, DEMANDS, OBLIGATIONS, DISPUTES, DAMAGES, LOSSES, EXPENSES, ATTORNEYS’ FEES, COSTS, LIABILITIES, AND ALL OTHER CLAIMS OF EVERY KIND AND NATURE WHATSOEVER, WHETHER KNOWN OR UNKNOWN, ANTICIPATED OR UNANTICIPATED, FORESEEN OR UNFORESEEN, ACCRUED OR UNACCRUED, DISCLOSED OR UNDISCLOSED, ARISING OUT OF OR RELATING TO LICENSEE’S USE OF THE SOFTWARE. NOTHING IN THIS SECTION SHOULD BE INTERPRETED TO REQUIRE LICENSEE TO INDEMNIFY LICENSOR, NOR REQUIRE LICENSOR TO INDEMNIFY LICENSEE.
+
+**[7.](#7) TERMINATION**
+
+[7.1.](#7.1) _Violations of Ethical Standards or Breaching Duty of Care_: If Licensee violates the Ethical Standards section or Licensee, or any other person or entity within the Supply Chain prior to a good or service reaching the Licensee, breaches its Duty of Care to Supply Chain Impacted Parties, Licensee must remedy the violation or harm caused by Licensee within 30 days of being notified of the violation or harm. If Licensee fails to remedy the violation or harm within 30 days, all rights in the Software granted to Licensee by License will be null and void as between Licensor and Licensee.
+
+[7.2.](#7.2) _Failure of Notice_: If any person or entity notifies Licensee in writing that Licensee has not complied with the Notice section of this License, Licensee can keep this License by taking all practical steps to comply within 30 days after the notice of noncompliance. If Licensee does not do so, Licensee’s License (and all rights licensed hereunder) will end immediately.
+
+[7.3.](#7.3) _Judicial Findings_: In the event Licensee is found by a civil, criminal, administrative, or other court of competent jurisdiction, or some other adjudicating body with legal authority, to have committed actions which are in violation of the Ethical Standards or Supply Chain Impacted Party sections of this License, all rights granted to Licensee by this License will terminate immediately.
+
+[7.4.](#7.4) _Patent Litigation_: If Licensee institutes patent litigation against any entity (including a cross-claim or counterclaim in a suit) alleging that the Software, all or part of the Software’s code, or a derivative work developed using the Software, including a portion of its code, constitutes direct or contributory patent infringement, then any patent license, along with all other rights, granted to Licensee under this License will terminate as of the date such litigation is filed.
+
+[7.5.](#7.5) _Additional Remedies_: Termination of the License by failing to remedy harms in no way prevents Licensor or Supply Chain Impacted Party from seeking appropriate remedies at law or in equity.
+
+**[8.](#8) MISCELLANEOUS:**
+
+[8.1.](#8.1) _Conditions_: Sections 3, 4.1, 5.1, 5.2, 7.1, 7.2, 7.3, and 7.4 are conditions of the rights granted to Licensee in the License.
+
+[8.2.](#8.2) _Equitable Relief_: Licensor and any Supply Chain Impacted Party shall be entitled to equitable relief, including injunctive relief or specific performance of the terms hereof, in addition to any other remedy to which they are entitled at law or in equity.
+
+[8.3.](#8.3) _Severability_: If any term or provision of this License is determined to be invalid, illegal, or unenforceable by a court of competent jurisdiction, any such determination of invalidity, illegality, or unenforceability shall not affect any other term or provision of this License or invalidate or render unenforceable such term or provision in any other jurisdiction. If the determination of invalidity, illegality, or unenforceability by a court of competent jurisdiction pertains to the terms or provisions contained in the Ethical Standards section of this License, all rights in the Software granted to Licensee shall be deemed null and void as between Licensor and Licensee.
+
+[8.4.](#8.4) _Section Titles_: Section titles are solely written for organizational purposes and should not be used to interpret the language within each section.
+
+[8.5.](#8.5) _Citations_: Citations are solely written to provide context for the source of the provisions in the Ethical Standards.
+
+[8.6.](#8.6) _Section Summaries_: Some sections have a brief _italicized description_ which is provided for the sole purpose of briefly describing the section and should not be used to interpret the terms of the License.
+
+[8.7.](#8.7) _Entire License_: This is the entire License between the Licensor and Licensee with respect to the claims released herein and that the consideration stated herein is the only consideration or compensation to be paid or exchanged between them for this License. This License cannot be modified or amended except in a writing signed by Licensor and Licensee.
+
+[8.8.](#8.8) _Successors and Assigns_: This License shall be binding upon and inure to the benefit of the Licensor’s and Licensee’s respective heirs, successors, and assigns.

cyxla-0.1.0/PKG-INFO

@@ -0,0 +1,235 @@
+Metadata-Version: 2.4
+Name: cyxla
+Version: 0.1.0
+Summary: A bridge from JAX to native code through the PJRT C API
+Author-Email: Nardi Lam <mail@nardilam.nl>
+License-File: LICENSE.md
+License-File: src/cyxla/include/xla/LICENSE
+Project-URL: Source, https://github.com/nardi/cyxla/tree/main/cyxla
+Project-URL: Documentation, https://nardi.github.io/cyxla/
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.10
+Requires-Dist: numpy>=1.26
+Provides-Extra: cpu
+Requires-Dist: cyxla-pjrt-cpu>=0.12; extra == "cpu"
+Provides-Extra: cuda12
+Requires-Dist: jax-cuda12-plugin[with-cuda]; extra == "cuda12"
+Provides-Extra: cuda13
+Requires-Dist: jax-cuda13-plugin[with-cuda]; extra == "cuda13"
+Provides-Extra: rocm7
+Requires-Dist: jax-rocm7-pjrt; extra == "rocm7"
+Provides-Extra: tpu
+Requires-Dist: libtpu; extra == "tpu"
+Description-Content-Type: text/markdown
+
+# cyxla
+
+cyxla lowers JAX functions to StableHLO, compiles them with a PJRT plugin, and runs them
+from Cython extension modules without the GIL. It suits native libraries that call JAX
+functions thousands of times per second, such as the SCIP callbacks of
+[JAX-CIP](https://github.com/nardi/jax-cip). The
+[documentation](https://nardi.github.io/cyxla/) has guides and the API reference.
+
+## Installation
+
+The package and a plugin install from wheels, and nothing downloads at run time.
+
+```bash
+pip install cyxla[cpu]
+```
+
+Each extra installs the plugin of one platform:
+
+| Extra | Platform | Plugin |
+| --- | --- | --- |
+| `cpu` | Linux on x86-64 and arm64, macOS on arm64, Windows on x86-64 | [cyxla-pjrt-cpu](../cyxla-pjrt-cpu) |
+| `cuda12`, `cuda13` | NVIDIA GPUs on Linux | JAX's `jax-cuda12-pjrt` or `jax-cuda13-pjrt`, with the CUDA libraries |
+| `rocm7` | AMD GPUs on Linux | JAX's `jax-rocm7-pjrt`, with ROCm from the system |
+| `tpu` | Cloud TPUs | `libtpu` |
+
+jaxlib links its own CPU client statically and ships no plugin library, which is why the
+CPU plugin is a separate package.
+
+## Example
+
+<!-- name: test_readme -->
+```python
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import cyxla as cx
+
+program = cx.lower(lambda x: jnp.sin(x) * 2, [jax.ShapeDtypeStruct((3,), jnp.float32)])
+executable, statistics = cx.compile_program(program)
+(result,) = executable.run([np.array([0.0, 1.0, 2.0])])
+```
+
+The linear algebra of `jnp.linalg` works the same way, as the section on custom calls
+explains:
+
+<!-- name: test_readme -->
+```python
+matrix = np.array([[4.0, 1.0], [1.0, 3.0]])
+program = cx.lower(jnp.linalg.cholesky, [jax.ShapeDtypeStruct((2, 2), jnp.float32)])
+executable, _ = cx.compile_program(program)
+(factor,) = executable.run([matrix])
+assert np.allclose(factor, np.linalg.cholesky(matrix))
+```
+
+The first function that needs a plugin loads the default one, so a program needs no
+setup call. JAX lowers to 32-bit types unless `jax.enable_x64` is on, as it does without
+cyxla. Arguments and results can be `bool`, 8 to 64-bit integers, `float16`, `bfloat16`,
+`float32`, `float64`, `complex64` and `complex128`.
+
+`Executable.run` returns arrays with the shapes and dtypes of the outputs, and converts
+arguments to the types that the program takes. It is for tests and interactive checks.
+Extension modules call `execute` from `cyxla/_runtime.pxd` instead, as the next section
+shows.
+
+## Extension modules
+
+An extension module cimports the runtime and runs programs with `execute`, which takes
+no GIL. On the CPU, inputs alias host arrays that the module fills in place, and outputs
+are read from the buffers' memory, so an execution copies nothing beyond the arrays that
+the caller passes. On other platforms every execution copies its inputs to the device
+and its outputs back.
+
+```python
+import cython
+from cython.cimports.cyxla._runtime import ABI_VERSION, Input, ProgramSpec, execute
+
+import cyxla
+
+cyxla.check_abi(ABI_VERSION)
+```
+
+`check_abi` rejects a module that was compiled against declarations of another cyxla
+version. Cython checks the signatures of shared functions itself, but not the layout of
+structs such as `ProgramSpec`.
+
+Host arrays for inputs come from `aligned_memory`, which takes a size in bytes, and go
+back with `free_aligned`. The pair exists because Windows frees aligned memory with a
+function other than `free`. Each `Input` names the PJRT element type of its array, such
+as `PJRT_Buffer_Type_F32` from `cyxla._pjrt`, and `ProgramSpec.output_bytes` holds the
+size of each output, which `Executable.output_bytes` returns.
+
+The build needs the PJRT header and the `.pxd` files, and `cyxla.get_include()` returns
+their directories. A Meson build reads them like this, and `tests/downstream` is a
+complete example.
+
+```meson
+cyxla_include = run_command(
+  python, '-c', 'import cyxla; print("\\n".join(cyxla.get_include()))', check: true,
+).stdout().strip().split('\n')
+# Pass cyxla_include[1] to Cython with -I, and cyxla_include[0] to the C compiler.
+```
+
+Downstream packages list `cyxla` in `build-system.requires` and pin it to a narrow
+range, because a new version may change the declarations.
+
+## Plugins
+
+A process loads one plugin. `cx.lower` picks a platform like JAX does, from the device
+that the arguments are committed to, then `jax.default_device`, then the loaded plugin,
+then `jax_platforms`, and last the CPU plugin if one is installed, then CUDA, ROCm and
+TPU in that order. The program records the platform, and compiling it loads the plugin
+for that platform. `cx.initialize` chooses a plugin or its options instead, and has to
+come before that first use. `cx.initialize(platform="cuda")` asks for a platform, and
+`cx.initialize(Path("libpjrt.so"), platform="cpu")` loads a library by path. `python -m
+cyxla` lists what it finds.
+
+Each plugin comes with the options that `PJRT_Client_Create` receives, and `options`
+adds to or replaces them:
+
+```python
+cx.initialize(platform="cuda", options={"memory_fraction": 0.5})
+```
+
+The CPU client runs synchronously on one device, so executions stay in the calling
+thread. The GPU clients do not preallocate most of the device memory, as JAX does by
+default. Before it loads JAX's CUDA plugin, cyxla loads the CUDA libraries of the
+`nvidia` wheels, the same way JAX does.
+
+Other packages can provide a plugin with an entry point in the `cyxla.plugins` group.
+The entry point names a module with `PLATFORM`, `CREATE_OPTIONS` and `library_path()`,
+as `cyxla_pjrt_cpu` has.
+
+## Custom calls
+
+Some JAX functions lower to custom calls, which are native functions that the program
+runs. On the CPU, `jnp.linalg` and `jax.scipy.linalg` call the LAPACK kernels that
+jaxlib ships. A PJRT plugin is a separate build of XLA that does not know them, so
+loading the plugin registers jaxlib's typed FFI handlers with it by name. Programs with
+`cholesky`, `solve`, `svd`, `eigh`, `qr`, `lu`, `eigvals` and the other linear algebra
+functions then compile and run like they do in JAX.
+
+A library with its own primitives registers a typed FFI handler with
+`jax.ffi.register_ffi_target`, and cyxla passes it on to the plugin. Importing cyxla
+mirrors every later registration, and loading the plugin registers the ones that JAX
+already holds, so the order of the imports does not matter. A library that does not use
+JAX calls `cx.register_ffi_target` with the same arguments. `jax.export` refuses custom
+calls it has no compatibility promise for, so `cx.lower` takes the names of the targets
+that the function may call.
+
+```python
+import my_library
+
+program = cx.lower(my_library.function, shapes, custom_calls=["my_target"])
+```
+
+The plugin accepts typed FFI handlers with an `execute` stage. It does not run the older
+untyped custom calls (`api_version=0`) or the Python callbacks of `jax.pure_callback`
+and `jax.experimental.io_callback`, and compiling a program with one raises an error
+that names the target.
+
+## Disk cache
+
+`compile_program` stores executables in the cache directory under a hash of their
+StableHLO, the compile options and the plugin. A later run of the same program loads the
+executable instead of compiling it. `store_lowered` and `load_lowered` keep lowered
+programs under a key from `cache_key`, which includes the versions of JAX, cyxla and any
+distributions that the caller names.
+
+## Environment variables
+
+| Variable | Meaning |
+| --- | --- |
+| `CYXLA_PJRT_PLUGIN` | Path of a plugin that replaces the installed ones |
+| `CYXLA_PLATFORM` | Platform to load when `initialize` gets none |
+| `CYXLA_CACHE_DIR` | Directory of the disk cache, `~/.cache/cyxla` if unset |
+| `CYXLA_MIRROR_JAX` | Set to `0` to keep cyxla from mirroring the FFI registrations of JAX |
+
+## Limitations
+
+- Buffers hold the dtypes listed above and no others, such as the 8-bit floats, 4-bit
+  integers and PRNG keys. `cx.lower` refuses a function that uses them. TPUs do not
+  support `float64`, so programs for them fail to compile.
+- Custom calls need a typed FFI handler. Untyped custom calls and Python callbacks do
+  not work.
+- The GPU and TPU paths are tested up to client creation. CI has no GPU, so execution on
+  one is untested.
+- On Windows only the CPU plugin exists, since JAX publishes no native GPU plugins
+  there.
+
+## Development
+
+```bash
+uv sync
+uv run pytest
+```
+
+`uv sync` builds the CPU plugin package from `../cyxla-pjrt-cpu`, which downloads its
+archive once. Development builds trace lines in the Cython module, so its coverage is
+measured too. The tests need 100% coverage.
+
+CI has no GPU, so [`notebooks/colab_gpu_tests.ipynb`](notebooks/colab_gpu_tests.ipynb)
+runs the CUDA plugin on a Colab GPU.
+
+[![Open In
+Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/nardi/cyxla/blob/main/cyxla/notebooks/colab_gpu_tests.ipynb)
+
+## License
+
+[Hippocratic License 3.0](LICENSE.md). The vendored PJRT header is under XLA's Apache
+2.0 license.

cyxla-0.1.0/README.md

@@ -0,0 +1,211 @@
+# cyxla
+
+cyxla lowers JAX functions to StableHLO, compiles them with a PJRT plugin, and runs them
+from Cython extension modules without the GIL. It suits native libraries that call JAX
+functions thousands of times per second, such as the SCIP callbacks of
+[JAX-CIP](https://github.com/nardi/jax-cip). The
+[documentation](https://nardi.github.io/cyxla/) has guides and the API reference.
+
+## Installation
+
+The package and a plugin install from wheels, and nothing downloads at run time.
+
+```bash
+pip install cyxla[cpu]
+```
+
+Each extra installs the plugin of one platform:
+
+| Extra | Platform | Plugin |
+| --- | --- | --- |
+| `cpu` | Linux on x86-64 and arm64, macOS on arm64, Windows on x86-64 | [cyxla-pjrt-cpu](../cyxla-pjrt-cpu) |
+| `cuda12`, `cuda13` | NVIDIA GPUs on Linux | JAX's `jax-cuda12-pjrt` or `jax-cuda13-pjrt`, with the CUDA libraries |
+| `rocm7` | AMD GPUs on Linux | JAX's `jax-rocm7-pjrt`, with ROCm from the system |
+| `tpu` | Cloud TPUs | `libtpu` |
+
+jaxlib links its own CPU client statically and ships no plugin library, which is why the
+CPU plugin is a separate package.
+
+## Example
+
+<!-- name: test_readme -->
+```python
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import cyxla as cx
+
+program = cx.lower(lambda x: jnp.sin(x) * 2, [jax.ShapeDtypeStruct((3,), jnp.float32)])
+executable, statistics = cx.compile_program(program)
+(result,) = executable.run([np.array([0.0, 1.0, 2.0])])
+```
+
+The linear algebra of `jnp.linalg` works the same way, as the section on custom calls
+explains:
+
+<!-- name: test_readme -->
+```python
+matrix = np.array([[4.0, 1.0], [1.0, 3.0]])
+program = cx.lower(jnp.linalg.cholesky, [jax.ShapeDtypeStruct((2, 2), jnp.float32)])
+executable, _ = cx.compile_program(program)
+(factor,) = executable.run([matrix])
+assert np.allclose(factor, np.linalg.cholesky(matrix))
+```
+
+The first function that needs a plugin loads the default one, so a program needs no
+setup call. JAX lowers to 32-bit types unless `jax.enable_x64` is on, as it does without
+cyxla. Arguments and results can be `bool`, 8 to 64-bit integers, `float16`, `bfloat16`,
+`float32`, `float64`, `complex64` and `complex128`.
+
+`Executable.run` returns arrays with the shapes and dtypes of the outputs, and converts
+arguments to the types that the program takes. It is for tests and interactive checks.
+Extension modules call `execute` from `cyxla/_runtime.pxd` instead, as the next section
+shows.
+
+## Extension modules
+
+An extension module cimports the runtime and runs programs with `execute`, which takes
+no GIL. On the CPU, inputs alias host arrays that the module fills in place, and outputs
+are read from the buffers' memory, so an execution copies nothing beyond the arrays that
+the caller passes. On other platforms every execution copies its inputs to the device
+and its outputs back.
+
+```python
+import cython
+from cython.cimports.cyxla._runtime import ABI_VERSION, Input, ProgramSpec, execute
+
+import cyxla
+
+cyxla.check_abi(ABI_VERSION)
+```
+
+`check_abi` rejects a module that was compiled against declarations of another cyxla
+version. Cython checks the signatures of shared functions itself, but not the layout of
+structs such as `ProgramSpec`.
+
+Host arrays for inputs come from `aligned_memory`, which takes a size in bytes, and go
+back with `free_aligned`. The pair exists because Windows frees aligned memory with a
+function other than `free`. Each `Input` names the PJRT element type of its array, such
+as `PJRT_Buffer_Type_F32` from `cyxla._pjrt`, and `ProgramSpec.output_bytes` holds the
+size of each output, which `Executable.output_bytes` returns.
+
+The build needs the PJRT header and the `.pxd` files, and `cyxla.get_include()` returns
+their directories. A Meson build reads them like this, and `tests/downstream` is a
+complete example.
+
+```meson
+cyxla_include = run_command(
+  python, '-c', 'import cyxla; print("\\n".join(cyxla.get_include()))', check: true,
+).stdout().strip().split('\n')
+# Pass cyxla_include[1] to Cython with -I, and cyxla_include[0] to the C compiler.
+```
+
+Downstream packages list `cyxla` in `build-system.requires` and pin it to a narrow
+range, because a new version may change the declarations.
+
+## Plugins
+
+A process loads one plugin. `cx.lower` picks a platform like JAX does, from the device
+that the arguments are committed to, then `jax.default_device`, then the loaded plugin,
+then `jax_platforms`, and last the CPU plugin if one is installed, then CUDA, ROCm and
+TPU in that order. The program records the platform, and compiling it loads the plugin
+for that platform. `cx.initialize` chooses a plugin or its options instead, and has to
+come before that first use. `cx.initialize(platform="cuda")` asks for a platform, and
+`cx.initialize(Path("libpjrt.so"), platform="cpu")` loads a library by path. `python -m
+cyxla` lists what it finds.
+
+Each plugin comes with the options that `PJRT_Client_Create` receives, and `options`
+adds to or replaces them:
+
+```python
+cx.initialize(platform="cuda", options={"memory_fraction": 0.5})
+```
+
+The CPU client runs synchronously on one device, so executions stay in the calling
+thread. The GPU clients do not preallocate most of the device memory, as JAX does by
+default. Before it loads JAX's CUDA plugin, cyxla loads the CUDA libraries of the
+`nvidia` wheels, the same way JAX does.
+
+Other packages can provide a plugin with an entry point in the `cyxla.plugins` group.
+The entry point names a module with `PLATFORM`, `CREATE_OPTIONS` and `library_path()`,
+as `cyxla_pjrt_cpu` has.
+
+## Custom calls
+
+Some JAX functions lower to custom calls, which are native functions that the program
+runs. On the CPU, `jnp.linalg` and `jax.scipy.linalg` call the LAPACK kernels that
+jaxlib ships. A PJRT plugin is a separate build of XLA that does not know them, so
+loading the plugin registers jaxlib's typed FFI handlers with it by name. Programs with
+`cholesky`, `solve`, `svd`, `eigh`, `qr`, `lu`, `eigvals` and the other linear algebra
+functions then compile and run like they do in JAX.
+
+A library with its own primitives registers a typed FFI handler with
+`jax.ffi.register_ffi_target`, and cyxla passes it on to the plugin. Importing cyxla
+mirrors every later registration, and loading the plugin registers the ones that JAX
+already holds, so the order of the imports does not matter. A library that does not use
+JAX calls `cx.register_ffi_target` with the same arguments. `jax.export` refuses custom
+calls it has no compatibility promise for, so `cx.lower` takes the names of the targets
+that the function may call.
+
+```python
+import my_library
+
+program = cx.lower(my_library.function, shapes, custom_calls=["my_target"])
+```
+
+The plugin accepts typed FFI handlers with an `execute` stage. It does not run the older
+untyped custom calls (`api_version=0`) or the Python callbacks of `jax.pure_callback`
+and `jax.experimental.io_callback`, and compiling a program with one raises an error
+that names the target.
+
+## Disk cache
+
+`compile_program` stores executables in the cache directory under a hash of their
+StableHLO, the compile options and the plugin. A later run of the same program loads the
+executable instead of compiling it. `store_lowered` and `load_lowered` keep lowered
+programs under a key from `cache_key`, which includes the versions of JAX, cyxla and any
+distributions that the caller names.
+
+## Environment variables
+
+| Variable | Meaning |
+| --- | --- |
+| `CYXLA_PJRT_PLUGIN` | Path of a plugin that replaces the installed ones |
+| `CYXLA_PLATFORM` | Platform to load when `initialize` gets none |
+| `CYXLA_CACHE_DIR` | Directory of the disk cache, `~/.cache/cyxla` if unset |
+| `CYXLA_MIRROR_JAX` | Set to `0` to keep cyxla from mirroring the FFI registrations of JAX |
+
+## Limitations
+
+- Buffers hold the dtypes listed above and no others, such as the 8-bit floats, 4-bit
+  integers and PRNG keys. `cx.lower` refuses a function that uses them. TPUs do not
+  support `float64`, so programs for them fail to compile.
+- Custom calls need a typed FFI handler. Untyped custom calls and Python callbacks do
+  not work.
+- The GPU and TPU paths are tested up to client creation. CI has no GPU, so execution on
+  one is untested.
+- On Windows only the CPU plugin exists, since JAX publishes no native GPU plugins
+  there.
+
+## Development
+
+```bash
+uv sync
+uv run pytest
+```
+
+`uv sync` builds the CPU plugin package from `../cyxla-pjrt-cpu`, which downloads its
+archive once. Development builds trace lines in the Cython module, so its coverage is
+measured too. The tests need 100% coverage.
+
+CI has no GPU, so [`notebooks/colab_gpu_tests.ipynb`](notebooks/colab_gpu_tests.ipynb)
+runs the CUDA plugin on a Colab GPU.
+
+[![Open In
+Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/nardi/cyxla/blob/main/cyxla/notebooks/colab_gpu_tests.ipynb)
+
+## License
+
+[Hippocratic License 3.0](LICENSE.md). The vendored PJRT header is under XLA's Apache
+2.0 license.

cyxla-0.1.0/docs/api/cache.md

@@ -0,0 +1,13 @@
+# Disk cache
+
+::: cyxla.cache_directory
+
+::: cyxla.cache_key
+
+::: cyxla.library_identity
+
+::: cyxla.store_lowered
+
+::: cyxla.load_lowered
+
+::: cyxla.CACHE_ENVIRONMENT_VARIABLE

cyxla-0.1.0/docs/api/custom-calls.md

@@ -0,0 +1,12 @@
+# Custom calls
+
+Registering the native functions that programs call. The [custom calls
+guide](../guide/custom-calls.md) explains when each is needed.
+
+::: cyxla.register_ffi_target
+
+::: cyxla.register_jax_targets
+
+::: cyxla.mirror_jax_registrations
+
+::: cyxla.MIRROR_ENVIRONMENT_VARIABLE

cyxla-0.1.0/docs/api/extension.md

@@ -0,0 +1,13 @@
+# Extension modules
+
+What a module that cimports the runtime needs at build and import time. The functions
+that the module calls, such as `execute`, are declared in `cyxla/_runtime.pxd` and are
+described in the [extension modules guide](../guide/extensions.md).
+
+::: cyxla.get_include
+
+::: cyxla.check_abi
+
+::: cyxla.ABI_VERSION
+
+::: cyxla.abi_version