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

This version

0.1.2

Jul 27, 2023

0.1.1

Jul 18, 2023

0.1.0

Jul 12, 2023

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.2.tar.gz (15.9 kB view hashes)

Uploaded Jul 27, 2023 Source

Built Distribution

jax_enums-0.1.2-py2.py3-none-any.whl (12.9 kB view hashes)

Uploaded Jul 27, 2023 Python 2 Python 3

Hashes for jax_enums-0.1.2.tar.gz

Hashes for jax_enums-0.1.2.tar.gz
Algorithm	Hash digest
SHA256	`75391a59e1f5be3a976715306182d22331d74c001ae3b580e95a88c03d58dabf`
MD5	`1d1230a926ef2557794d46818e352d90`
BLAKE2b-256	`c97e86dadf746bba5e0c7cfd08db797dfd9e52c9bbe904ee74998727cac17381`

Hashes for jax_enums-0.1.2-py2.py3-none-any.whl

Hashes for jax_enums-0.1.2-py2.py3-none-any.whl
Algorithm	Hash digest
SHA256	`e9c9374789f9adb5874013c80e2be411af0ef1442cbc3c3c0bb0ac89de7f2f2c`
MD5	`44be5884e9bd2994463b5ce52df2a073`
BLAKE2b-256	`c9744cc28410c9f025ccb575123f35941f2812e67afd975f42b99682152e939e`