{"id":18417000,"url":"https://github.com/young-geng/mlxu","last_synced_at":"2025-07-17T20:37:50.331Z","repository":{"id":65556426,"uuid":"594620891","full_name":"young-geng/mlxu","owner":"young-geng","description":"Machine Learning eXperiment Utilities","archived":false,"fork":false,"pushed_at":"2025-06-27T07:52:01.000Z","size":46,"stargazers_count":46,"open_issues_count":3,"forks_count":8,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-06-27T08:15:25.515Z","etag":null,"topics":["jax","machine-learning","python","wandb"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/young-geng.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-01-29T05:21:01.000Z","updated_at":"2025-06-27T07:52:04.000Z","dependencies_parsed_at":"2024-06-11T06:24:22.758Z","dependency_job_id":"31be6160-3fac-4714-ad1d-e9667c453c50","html_url":"https://github.com/young-geng/mlxu","commit_stats":{"total_commits":16,"total_committers":1,"mean_commits":16.0,"dds":0.0,"last_synced_commit":"395722c0b1913a7db492796acaba365c26eb7fbc"},"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/young-geng/mlxu","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/young-geng%2Fmlxu","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/young-geng%2Fmlxu/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/young-geng%2Fmlxu/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/young-geng%2Fmlxu/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/young-geng","download_url":"https://codeload.github.com/young-geng/mlxu/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/young-geng%2Fmlxu/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":262762468,"owners_count":23360330,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["jax","machine-learning","python","wandb"],"created_at":"2024-11-06T04:07:53.081Z","updated_at":"2025-06-30T11:06:07.315Z","avatar_url":"https://github.com/young-geng.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"# MLXU: Machine Learning eXperiment Utilities\nThis library provide a collection of utilities for machine learning experiments.\nMLXU is a thin wrapper on top of [absl-py](https://github.com/abseil/abseil-py),\n[ml_collections](https://github.com/google/ml_collections) and\n[wandb](https://github.com/wandb/wandb). It also provides some convenient JAX\nutils.\n\n\nThis library includes the following modules:\n * [config](mlxu/config.py) Experiment configuration and command line flags utils\n * [logging](mlxu/logging.py) W\u0026B logging utils\n * [jax_utils](mlxu/jax_utils.py) JAX specific utils\n\n\n# Installation\nMLXU can be installed via pip. To install from PYPI\n```shell\npip install mlxu\n```\n\nTo install the latest version from GitHub\n```shell\npip install git+https://github.com/young-geng/mlxu.git\n```\n\n\n# Examples\nHere are some examples for the utilities provide in MLXU\n\n## Command Line Flags and Logging\nMLXU provides convenient wrappers around absl-py and wandb to make command line\narg parsing and logging easy.\n```python\nimport mlxu\n\n\nclass ConfigurableModule(object):\n    # Define a configurable module with a default configuration. This module\n    # can be directly configured from the command line when plugged into\n    # the FLAGS.\n\n    @staticmethod\n    def get_default_config(updates=None):\n        config = mlxu.config_dict()\n        config.integer_value = 10\n        config.float_value = 1.0\n        config.string_value = 'hello'\n        config.boolean_value = True\n        return mlxu.update_config_dict(config, updates)\n\n    def __init__(self, config):\n        self.config = self.get_default_config(config)\n\n\n# Define absl command line flags in one function, with automatic type inference.\nFLAGS, FLAGS_DEF = mlxu.define_flags_with_default(\n    name='example_experiment',          # string flag\n    seed=42,                            # integer flag\n    learning_rate=1e-3,                 # floating point flag\n    use_mlxu=True,                      # boolean flag\n    num_gpus=int,                       # integer flag without default value\n    weight_decay=float,                 # floating point flag without default value\n    save_checkpoints=bool,              # boolean flag without default value\n    epochs=(10, 'Number of epochs'),    # we can also specify help strings\n    network_architecture=mlxu.config_dict(\n        activation='relu',\n        hidden_dim=128,\n        hidden_layers=5,\n    ),                                  # nested ml_collections config_dict\n    configurable_module=ConfigurableModule.get_default_config(),  # nested custom config_dict\n    logger=mlxu.WandBLogger.get_default_config(),  # logger configuration\n)\n\n\ndef main(argv):\n    # Print the command line flags\n    mlxu.print_flags(FLAGS, FLAGS_DEF)\n\n    # Access the flags\n    name = FLAGS.name\n    seed = FLAGS.seed\n\n    # Access nested flags\n    activation = FLAGS.network_architecture.activation\n    hidden_dim = FLAGS.network_architecture.hidden_dim\n\n    configurable_module = ConfigurableModule(FLAGS.configurable_module)\n\n    # Create logger and log metrics\n    logger = mlxu.WandBLogger(FLAGS.logger, mlxu.get_user_flags(FLAGS, FLAGS_DEF))\n    logger.log({'step': 1, 'loss': 10.5})\n    logger.save_pickle([1, 2, 4, 5], 'checkpoint.pkl')\n\n\n# Run the main function\nif __name__ == \"__main__\":\n    mlxu.run(main)\n```\n\nThe flags can be passed in via command line arguments:\n```shell\npython examples/cli_logging.py \\\n    --name='example' \\\n    --seed=24 \\\n    --learning_rate=1.0 \\\n    --use_mlxu=True \\\n    --network_architecture.activation='gelu' \\\n    --network_architecture.hidden_dim=126 \\\n    --network_architecture.hidden_layers=2 \\\n    --configurable_module.integer_value=20 \\\n    --configurable_module.float_value=2.0 \\\n    --logger.online=True \\\n    --logger.project='mlxu_example'\n```\n\nSpecifically, the `logger.online` option controls whether the logger will upload\nthe data to W\u0026B, and the `logger.project` option controls the name of the W\u0026B\nproject.\n\n## JAX Random Number Generator\nMLXU also provides convenient wrapper around JAX's random number generators\nto make it much easier to use\n```python\nimport jax\nimport jax.numpy as jnp\nimport mlxu\nimport mlxu.jax_utils as jax_utils\n\n\n@jax.jit\ndef sum_of_random_uniform(rng_key):\n    # Capture RNG key to create a stateful rng key generator.\n    # As long as RNGGenerator object is not pass through the function\n    # boundary, the function is still pure and jittable.\n    # RNGGenerator object also supports the same tuple and dictionary usage like\n    # the jax_utils.next_rng function.\n    rng_generator = jax_utils.RNGGenerator(rng_key)\n    output = jnp.zeros((2, 2))\n    for i in range(4):\n        # Each call returns a new key, altering the internal state of rng_generator\n        output += jax.random.uniform(rng_generator(), (2, 2))\n\n    return output\n\n\ndef main(argv):\n    # Setup global rng generator\n    jax_utils.init_rng(42)\n\n    # Get an rng key\n    rng_key = jax_utils.next_rng()\n    print(rng_key)\n\n    # Get a new rng key, this key should be different from the previous one\n    rng_key = jax_utils.next_rng()\n    print(rng_key)\n\n    # You can also get a tuple of N rng keys\n    k1, k2, k3 = jax_utils.next_rng(3)\n    print(k1, ', ', k2, ', ', k3)\n\n    # Dictionary of keys is also supported\n    rng_key_dict = jax_utils.next_rng(['k1', 'k2'])\n    print(rng_key_dict)\n\n    # Call a jitted function that makes use of stateful RNGGenerator object internally\n    x = sum_of_random_uniform(jax_utils.next_rng())\n    print(x)\n\n\nif __name__ == \"__main__\":\n    mlxu.run(main)\n```","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fyoung-geng%2Fmlxu","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fyoung-geng%2Fmlxu","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fyoung-geng%2Fmlxu/lists"}