JAX-compatible Enumerations.
Project description
JAX_ENUMS
Installation | Examples | Cite
A Jax-compatible enumerable.
Installation
pip install jax_enums
Example
class Foo(Enumerable):
BAR = 0
BAZ = 1
def f(array: jax.Array, enumerable: Enum) -> jax.Array:
return array[enumerable.value]
array = jnp.zeros((2, 2))
enumerable = Foo.BAR
f(array, enumerable)
jax.jit(f)(array, enumerable)
Cite
@misc{pignatelli2023jax_enums,
author = {Pignatelli, Eduardo},
title = {JAX_ENUMS: JAX-compatible enumerations},
year = {2023},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/epignatelli/jax_enums}}
}
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
jax_enums-0.1.1.tar.gz
(10.9 kB
view hashes)
Built Distribution
Close
Hashes for jax_enums-0.1.1-py2.py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 2d85ca4f2a29993b7b0b13583c36ee9da5861385f69e842dafba7c59120aa9e6 |
|
MD5 | a43d181069c12a42fef908a07efe5e1a |
|
BLAKE2b-256 | d4ce0441546cffc1bbf28fafc07ccd0c4f1c40ce03d40d32736299264a314d25 |