PyPI page
Author:
Mahdi Shafiei
License:
MIT License
Copyright (c) 2026 Mahdi Shafiei
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated docu...
Summary:
JAX benchmarking, profiling and evaluation metrics for Flax NNX: a metric registry across regression, classification, calibration, uncertainty, forecasting, generative, image, text, audio, graph and fairness domains, XLA FLOP counting, roofline analysis, GPU and energy monitoring, regression detection, publication exports, W&B and MLflow
Latest version:
0.1.14
Required dependencies:
click
|
flax
|
jax
|
jaxlib
|
jaxtyping
|
lazy-loader
|
numpy
|
psutil
|
pydantic
|
scipy
|
substrax
Optional dependencies:
bandit
|
beartype
|
build
|
codecarbon
|
coverage
|
griffe
|
hippogriffe
|
hypothesis
|
import-linter
|
interrogate
|
ipykernel
|
jax
|
jax-metal
|
jupytext
|
matplotlib
|
mkdocs
|
mkdocs-include-exclude-files
|
mkdocs-ipynb
|
mkdocs-jupyter
|
mkdocs-material
|
mkdocstrings
|
mkdocstrings-python
|
mlflow
|
nvidia-cublas-cu12
|
nvidia-cuda-cupti-cu12
|
nvidia-cuda-nvcc-cu12
|
nvidia-cuda-nvrtc-cu12
|
nvidia-cuda-runtime-cu12
|
nvidia-cudnn-cu12
|
nvidia-cufft-cu12
|
nvidia-curand-cu12
|
nvidia-cusolver-cu12
|
nvidia-cusparse-cu12
|
nvidia-ml-py
|
nvidia-nccl-cu12
|
nvidia-nvjitlink-cu12
|
nvidia-nvtx-cu12
|
pip-audit
|
pre-commit
|
pydeps
|
pydoclint
|
pygments
|
pymdown-extensions
|
pyright
|
pytest
|
pytest-asyncio
|
pytest-benchmark
|
pytest-codspeed
|
pytest-cov
|
pytest-env
|
pytest-json-report
|
pytest-randomly
|
pytest-timeout
|
pytest-xdist
|
radon
|
ruff
|
ruptures
|
scikit-learn
|
scipy-stubs
|
substrax
|
twine
|
wandb
Downloads last day:
113
Downloads last week:
1,481
Downloads last month:
13,515
7-day smoothing