adv-jax-math

adv-jax-math provides batching and sharding utilities for vectorization and differentiation in JAX.

Installation

adv-jax-math is installable with pip:

pip install adv-jax-math

For full details, see the API documentation.

Licensing

Copyright (C) 2024 Kaya Unalmis.

adv-jax-math is licensed under the GNU Lesser General Public License, version 3.

Public API

• API Documentation
  • Batched maps
    • adv_jax_math.batch_jacfwd
    • adv_jax_math.batch_jacrev
    • adv_jax_math.batch_map
    • adv_jax_math.batch_vectorize
    • adv_jax_math.batch_vmap
  • Sparse differentiation
    • adv_jax_math.sparse_pullback
    • adv_jax_math.sparse_pullback_map

Indices and tables

• Index

• Module Index

• Search Page