Yu-Hang "Maxin" Tang

Install non-Python dependencies with `dnf`

5

Hi there, This is related to the 'installation' checklist item of the JOSS review: https://github.com/openjournals/joss-reviews/issues/4149. Could you please also provide the external package names under the DNF/RPM system? This could...

Add Kubernetes Jobset support to `jax.distributed.initialize`

1

This should allow JAX to automatically bootstrap distributed computations when running inside a [Kubernetes jobset](https://github.com/kubernetes-sigs/jobset).

Allow using `python -m grpc_tools.protoc` in lieu of `protoc`

2

Thank you for creating this valuable package! Currently, `protol`'s `protoc` command relies on system calls to the `protoc` executable. Since `protoc` is not available as a PyPI package, this creates...

Use lambda instead of functools.partial to create single-host config maker

4

Fixes #1151 An alternative solution is to use `functools.wraps` but that requires a custom config validator: ``` axlearn.common.config.InvalidConfigValueError: Invalid config value type for value "functools.partial(, 'fuji-3B-v3-flash')". Consider registering a custom...

I encountered the following error when trying to run the Fuji 3B model locally: ``` Traceback (most recent call last): File "fuji.py", line 339, in run_main() File "fuji.py", line 219,...

Move Gemma's stacked_params_mapping to class scope

5

This aligns Gemma with the implementations of Gemma2/3, where `packed_modules_mapping` is defined as a class attribute. This allows us to query the stacking specifications from outside.