# nshutils
`nshutils` is a collection of utility functions and classes that I've found useful in my day-to-day work as an ML researcher. This library includes utilities for typechecking, logging, and saving/loading activations from neural networks.
## Installation
To install `nshutils`, simply run:
```bash
pip install nshutils
```
## Features
### Typechecking
`nshutils` provides a simple way to typecheck your code using the [`jaxtyping`](https://github.com/patrick-kidger/jaxtyping) library. Simply call `typecheck_this_module()` at the top of your module (i.e., in the root `__init__.py` file) to enable typechecking for the entire module:
```python
from nshutils.typecheck import typecheck_this_module
typecheck_this_module()
```
You can also use the `tassert` function to assert that a value is of a certain type:
```python
import nshutils.typecheck as tc
def my_function(x: tc.Float[torch.Tensor, "bsz seq len"]) -> tc.Float[torch.Tensor, "bsz seq len"]:
tc.tassert(tc.Float[torch.Tensor, "bsz seq len"], x)
...
```
### Logging
`nshutils` provides a simple way to configure logging for your project. Simply call one of the logging setup functions:
```python
from nshutils.logging import init_python_logging
init_python_logging()
```
This will configure logging to use pretty formatting for PyTorch tensors and numpy arrays (inspired by and/or utilizing [`lovely-numpy`](https://github.com/xl0/lovely-numpy) and [`lovely-tensors`](https://github.com/xl0/lovely-tensors)), and will also enable rich logging if the `rich` library is installed.
### Activation Saving/Loading
`nshutils` provides a simple way to save and load activations from neural networks. To save activations, use the `ActSave` object:
```python
from nshutils import ActSave
def my_model_forward(x):
...
# Save activations to "{save_dir}/encoder.activations/{idx}.npy"
ActSave({"encoder.activations": x})
# Equivalent to the above
with ActSave.context("encoder"):
ActSave(activations=x)
...
ActSave.enable(save_dir="path/to/activations")
x = torch.randn(...)
my_model_forward(x)
# Activations are saved to disk under the "path/to/activations" directory
```
This will save the `x` tensor to disk under the `encoder` prefix.
To load activations, use the `ActLoad` class:
```python
from nshutils import ActLoad
act_load = ActLoad.from_latest_version("path/to/activations")
encoder_acts = act_load["encoder"]
for act in encoder_acts:
print(act.shape)
```
This will load all of the activations saved under the `encoder` prefix.
### Other Utilities
`nshutils` also provides a few other utility functions/classes:
- `snoop`: A simple way to debug your code using the `pysnooper` library, based on the [`torchsnooper`](https://github.com/zasdfgbnm/TorchSnooper) library.
- `apply_to_collection`: Recursively apply a function to all elements of a collection that match a certain type, taken from the [`pytorch-lightning`](https://github.com/Lightning-AI/pytorch-lightning) library.
## Contributing
Contributions to `nshutils` are welcome! Please open an issue or submit a pull request on the [GitHub repository](https://github.com/nimashoghi/nshutils).
## License
`nshutils` is released under the MIT License. See the `LICENSE` file for more details.
Raw data
{
"_id": null,
"home_page": "https://github.com/nimashoghi/nshutils",
"name": "nshutils",
"maintainer": null,
"docs_url": null,
"requires_python": "<4.0,>=3.10",
"maintainer_email": null,
"keywords": null,
"author": "Nima Shoghi",
"author_email": "nimashoghi@gmail.com",
"download_url": "https://files.pythonhosted.org/packages/62/b6/fa37ddcf1ecbf6a3be8695eb21c4deba41b6ae300d3545853f9375f0c71b/nshutils-0.37.5.tar.gz",
"platform": null,
"description": "# nshutils\n\n`nshutils` is a collection of utility functions and classes that I've found useful in my day-to-day work as an ML researcher. This library includes utilities for typechecking, logging, and saving/loading activations from neural networks.\n\n## Installation\n\nTo install `nshutils`, simply run:\n\n```bash\npip install nshutils\n```\n\n## Features\n\n### Typechecking\n\n`nshutils` provides a simple way to typecheck your code using the [`jaxtyping`](https://github.com/patrick-kidger/jaxtyping) library. Simply call `typecheck_this_module()` at the top of your module (i.e., in the root `__init__.py` file) to enable typechecking for the entire module:\n\n```python\nfrom nshutils.typecheck import typecheck_this_module\n\ntypecheck_this_module()\n```\n\nYou can also use the `tassert` function to assert that a value is of a certain type:\n\n```python\nimport nshutils.typecheck as tc\n\ndef my_function(x: tc.Float[torch.Tensor, \"bsz seq len\"]) -> tc.Float[torch.Tensor, \"bsz seq len\"]:\n tc.tassert(tc.Float[torch.Tensor, \"bsz seq len\"], x)\n ...\n```\n\n### Logging\n\n`nshutils` provides a simple way to configure logging for your project. Simply call one of the logging setup functions:\n\n```python\nfrom nshutils.logging import init_python_logging\n\ninit_python_logging()\n```\n\nThis will configure logging to use pretty formatting for PyTorch tensors and numpy arrays (inspired by and/or utilizing [`lovely-numpy`](https://github.com/xl0/lovely-numpy) and [`lovely-tensors`](https://github.com/xl0/lovely-tensors)), and will also enable rich logging if the `rich` library is installed.\n\n### Activation Saving/Loading\n\n`nshutils` provides a simple way to save and load activations from neural networks. To save activations, use the `ActSave` object:\n\n```python\nfrom nshutils import ActSave\n\ndef my_model_forward(x):\n ...\n # Save activations to \"{save_dir}/encoder.activations/{idx}.npy\"\n ActSave({\"encoder.activations\": x})\n\n # Equivalent to the above\n with ActSave.context(\"encoder\"):\n ActSave(activations=x)\n ...\n\nActSave.enable(save_dir=\"path/to/activations\")\nx = torch.randn(...)\nmy_model_forward(x)\n# Activations are saved to disk under the \"path/to/activations\" directory\n```\n\nThis will save the `x` tensor to disk under the `encoder` prefix.\n\nTo load activations, use the `ActLoad` class:\n\n```python\nfrom nshutils import ActLoad\n\nact_load = ActLoad.from_latest_version(\"path/to/activations\")\nencoder_acts = act_load[\"encoder\"]\n\nfor act in encoder_acts:\n print(act.shape)\n```\n\nThis will load all of the activations saved under the `encoder` prefix.\n\n### Other Utilities\n\n`nshutils` also provides a few other utility functions/classes:\n\n- `snoop`: A simple way to debug your code using the `pysnooper` library, based on the [`torchsnooper`](https://github.com/zasdfgbnm/TorchSnooper) library.\n- `apply_to_collection`: Recursively apply a function to all elements of a collection that match a certain type, taken from the [`pytorch-lightning`](https://github.com/Lightning-AI/pytorch-lightning) library.\n\n## Contributing\n\nContributions to `nshutils` are welcome! Please open an issue or submit a pull request on the [GitHub repository](https://github.com/nimashoghi/nshutils).\n\n## License\n\n`nshutils` is released under the MIT License. See the `LICENSE` file for more details.\n",
"bugtrack_url": null,
"license": null,
"summary": null,
"version": "0.37.5",
"project_urls": {
"Homepage": "https://github.com/nimashoghi/nshutils"
},
"split_keywords": [],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "41b9b1e7e98bab9c91a8bcab23c4b709aa0221ca019a3f634cbfb63ad2e850a9",
"md5": "0b65430710f8f927039194bfb5c32601",
"sha256": "bb7ee194a7130aec1ae6695c06331ba34b0099641f18006bdcfb3b940f6c51d9"
},
"downloads": -1,
"filename": "nshutils-0.37.5-py3-none-any.whl",
"has_sig": false,
"md5_digest": "0b65430710f8f927039194bfb5c32601",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": "<4.0,>=3.10",
"size": 72388,
"upload_time": "2025-07-23T20:20:23",
"upload_time_iso_8601": "2025-07-23T20:20:23.500497Z",
"url": "https://files.pythonhosted.org/packages/41/b9/b1e7e98bab9c91a8bcab23c4b709aa0221ca019a3f634cbfb63ad2e850a9/nshutils-0.37.5-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "62b6fa37ddcf1ecbf6a3be8695eb21c4deba41b6ae300d3545853f9375f0c71b",
"md5": "e5fbc3f53bd415cef6ba0619faf381c9",
"sha256": "85be746790f8bbb194a45478afaed5f8cbc42714694c86c1e5d62471f6463951"
},
"downloads": -1,
"filename": "nshutils-0.37.5.tar.gz",
"has_sig": false,
"md5_digest": "e5fbc3f53bd415cef6ba0619faf381c9",
"packagetype": "sdist",
"python_version": "source",
"requires_python": "<4.0,>=3.10",
"size": 59499,
"upload_time": "2025-07-23T20:20:24",
"upload_time_iso_8601": "2025-07-23T20:20:24.666488Z",
"url": "https://files.pythonhosted.org/packages/62/b6/fa37ddcf1ecbf6a3be8695eb21c4deba41b6ae300d3545853f9375f0c71b/nshutils-0.37.5.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-07-23 20:20:24",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "nimashoghi",
"github_project": "nshutils",
"travis_ci": false,
"coveralls": false,
"github_actions": false,
"lcname": "nshutils"
}
JSON
