# FJUtils
a package for custom Jax Flax Functions and Utils
Welcome to FJUtils - A collection of useful functions and utilities for Flax and JAX!
Some Parts of code if from EasyLM <3
## Overview
FJUtils is a collection of functions and utilities that can help with various tasks when using Flax and JAX. It includes
checkpoint savers, partitioning tools, and other helpful functions.
The goal of FJUtils is to make your life easier when working with Flax and JAX. Whether you are training a new model,
fine-tuning an existing one, or just exploring the capabilities of these powerful frameworks, FJUtils has something to
offer.
## Features
Here are some of the features included in FJUtils:
Checkpoint saver: This tool provides an easy way to save and restore checkpoints during training. You can specify how
often to save checkpoints, where to store them, and more.
Partitioning tools: FJUtils includes several tools for partitioning data across multiple devices or nodes. These tools
can help you optimize the performance of your models on clusters or distributed systems.
Other utilities: FJUtils includes a variety of other helpful functions and utilities and more.
## Getting Started
To get started with FJUtils, simply install the package using pip:
```shell
pip install fjutils
```
Once installed, you can import the package and start using its functions and utilities. For example, here's how you can
use the checkpoint saver for loading models like :
```python
from fjutils import StreamingCheckpointer
ckpt = StreamingCheckpointer.load_trainstate_checkpoint('params::<path to model>')
```
or simply getting an optimizer for example adafactor with cosine scheduler :
```python
from jax import numpy as jnp
from fjutils.optimizers import get_adafactor_with_cosine_scheduler
optimizer, scheduler = get_adafactor_with_cosine_scheduler(
steps=5000,
learning_rate=5e-5,
weight_decay=1e-1,
min_dim_size_to_factor=128,
decay_rate=0.8,
decay_offset=0,
multiply_by_parameter_scale=True,
clipping_threshold=1.0,
momentum=None,
dtype_momentum=jnp.float32,
weight_decay_rate=None,
eps=1e-30,
factored=True,
weight_decay_mask=None,
)
```
or getting adamw with linear scheduler:
```python
from fjutils.optimizers import get_adamw_with_linear_scheduler
optimizer, scheduler = get_adamw_with_linear_scheduler(
steps=5000,
learning_rate_start=5e-5,
learning_rate_end=1e-5,
b1=0.9,
b2=0.999,
eps=1e-8,
eps_root=0.0,
weight_decay=1e-1,
mu_dtype=None,
)
```
## Documentation
Documentations are available [here](https://erfanzar.github.io/FJUtils/docs)
## Contributing
FJUtils is an open-source project, and contributions are always welcome! If you have a feature request, bug report, or
just want to help out with development, please check out our GitHub repository and feel free to submit a pull request or
open an issue.
Thank you for using FJUtils, and happy training!
## Credits
- some of the jax utilities are from EasyLM by [Young-Geng](https://github.com/young-geng)
Raw data
{
"_id": null,
"home_page": "https://github.com/erfanzar/",
"name": "FJUtils",
"maintainer": "",
"docs_url": null,
"requires_python": ">=3.7",
"maintainer_email": "",
"keywords": "",
"author": "Erfan Zare Chavoshi",
"author_email": "erfanzare82@yahoo.com",
"download_url": "https://files.pythonhosted.org/packages/3b/b9/d801aa2a8d42185a2e02865bf57f2ea4ca8fd258cc15a6eb786d1236f5ea/FJUtils-0.0.21.tar.gz",
"platform": null,
"description": "# FJUtils\r\n\r\na package for custom Jax Flax Functions and Utils\r\nWelcome to FJUtils - A collection of useful functions and utilities for Flax and JAX!\r\nSome Parts of code if from EasyLM <3\r\n\r\n## Overview\r\n\r\nFJUtils is a collection of functions and utilities that can help with various tasks when using Flax and JAX. It includes\r\ncheckpoint savers, partitioning tools, and other helpful functions.\r\nThe goal of FJUtils is to make your life easier when working with Flax and JAX. Whether you are training a new model,\r\nfine-tuning an existing one, or just exploring the capabilities of these powerful frameworks, FJUtils has something to\r\noffer.\r\n\r\n## Features\r\n\r\nHere are some of the features included in FJUtils:\r\n\r\nCheckpoint saver: This tool provides an easy way to save and restore checkpoints during training. You can specify how\r\noften to save checkpoints, where to store them, and more.\r\n\r\nPartitioning tools: FJUtils includes several tools for partitioning data across multiple devices or nodes. These tools\r\ncan help you optimize the performance of your models on clusters or distributed systems.\r\n\r\nOther utilities: FJUtils includes a variety of other helpful functions and utilities and more.\r\n\r\n## Getting Started\r\n\r\nTo get started with FJUtils, simply install the package using pip:\r\n\r\n```shell\r\npip install fjutils\r\n```\r\n\r\nOnce installed, you can import the package and start using its functions and utilities. For example, here's how you can\r\nuse the checkpoint saver for loading models like :\r\n\r\n```python\r\nfrom fjutils import StreamingCheckpointer\r\n\r\nckpt = StreamingCheckpointer.load_trainstate_checkpoint('params::<path to model>')\r\n\r\n```\r\n\r\nor simply getting an optimizer for example adafactor with cosine scheduler :\r\n\r\n```python\r\nfrom jax import numpy as jnp\r\nfrom fjutils.optimizers import get_adafactor_with_cosine_scheduler\r\n\r\noptimizer, scheduler = get_adafactor_with_cosine_scheduler(\r\n steps=5000,\r\n learning_rate=5e-5,\r\n weight_decay=1e-1,\r\n min_dim_size_to_factor=128,\r\n decay_rate=0.8,\r\n decay_offset=0,\r\n multiply_by_parameter_scale=True,\r\n clipping_threshold=1.0,\r\n momentum=None,\r\n dtype_momentum=jnp.float32,\r\n weight_decay_rate=None,\r\n eps=1e-30,\r\n factored=True,\r\n weight_decay_mask=None,\r\n)\r\n\r\n```\r\n\r\nor getting adamw with linear scheduler:\r\n\r\n```python\r\nfrom fjutils.optimizers import get_adamw_with_linear_scheduler\r\n\r\noptimizer, scheduler = get_adamw_with_linear_scheduler(\r\n steps=5000,\r\n learning_rate_start=5e-5,\r\n learning_rate_end=1e-5,\r\n b1=0.9,\r\n b2=0.999,\r\n eps=1e-8,\r\n eps_root=0.0,\r\n weight_decay=1e-1,\r\n mu_dtype=None,\r\n)\r\n\r\n```\r\n\r\n## Documentation\r\n\r\nDocumentations are available [here](https://erfanzar.github.io/FJUtils/docs)\r\n\r\n## Contributing\r\n\r\nFJUtils is an open-source project, and contributions are always welcome! If you have a feature request, bug report, or\r\njust want to help out with development, please check out our GitHub repository and feel free to submit a pull request or\r\nopen an issue.\r\n\r\nThank you for using FJUtils, and happy training!\r\n\r\n## Credits\r\n\r\n- some of the jax utilities are from EasyLM by [Young-Geng](https://github.com/young-geng)\r\n",
"bugtrack_url": null,
"license": "Apache License 2.0",
"summary": "",
"version": "0.0.21",
"project_urls": {
"Homepage": "https://github.com/erfanzar/"
},
"split_keywords": [],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "5b9842ad137f4b7e0678d420cf6743ec19ef461a7c0ee122e12bdd2e717b7bc8",
"md5": "f3dde35ba44639f5bcd3b46cb06e7861",
"sha256": "28d6129bf086d7ec9bc0dd8ddf2442b88befe00bfd22bded0cfbdb6c9acaf225"
},
"downloads": -1,
"filename": "FJUtils-0.0.21-py3-none-any.whl",
"has_sig": false,
"md5_digest": "f3dde35ba44639f5bcd3b46cb06e7861",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.7",
"size": 27711,
"upload_time": "2023-10-20T10:04:57",
"upload_time_iso_8601": "2023-10-20T10:04:57.379744Z",
"url": "https://files.pythonhosted.org/packages/5b/98/42ad137f4b7e0678d420cf6743ec19ef461a7c0ee122e12bdd2e717b7bc8/FJUtils-0.0.21-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "3bb9d801aa2a8d42185a2e02865bf57f2ea4ca8fd258cc15a6eb786d1236f5ea",
"md5": "65c4af86366406e84e4750d7216ba9b0",
"sha256": "051aec7c639913a3008015fcec7cde33bb67283fcc49c01360960a76a1f6c91e"
},
"downloads": -1,
"filename": "FJUtils-0.0.21.tar.gz",
"has_sig": false,
"md5_digest": "65c4af86366406e84e4750d7216ba9b0",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.7",
"size": 25689,
"upload_time": "2023-10-20T10:04:59",
"upload_time_iso_8601": "2023-10-20T10:04:59.656488Z",
"url": "https://files.pythonhosted.org/packages/3b/b9/d801aa2a8d42185a2e02865bf57f2ea4ca8fd258cc15a6eb786d1236f5ea/FJUtils-0.0.21.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2023-10-20 10:04:59",
"github": false,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"lcname": "fjutils"
}
JSON