orbax

orbax copied to clipboard

Use Checkpoint Manager to benchmark in Orbax V0.

copybara-service[bot]

Internal change

1

Internal change

copybara-service[bot]

Restore from single slice

1

1. Note that I'm still using an old version (`v0.11.15`) of `slice_devices` method. (This means that it will return the devices from single `slice` instead of single `replica`) 2. The...

findmyway

findmyway

Swap order of expected value and actual value when using [assertProto2Equal](https://source.corp.google.com/piper///depot/google3/net/proto2/contrib/pyutil/compare.py;rcl=781699748;l=153) to be consistent with definition.

1

Swap order of expected value and actual value when using [assertProto2Equal](https://source.corp.google.com/piper///depot/google3/net/proto2/contrib/pyutil/compare.py;rcl=781699748;l=153) to be consistent with definition.

copybara-service[bot]

Updated hns enabled check

deepikarajani24

Add DescriptorHandler test. Also fix is_abstract_handleable method.

copybara-service[bot]

Add tests for ProtoHandler.

copybara-service[bot]

[OBM-Juno] Capture jax mesh in Genmedia exports

1

[OBM-Juno] Capture jax mesh in Genmedia exports

copybara-service[bot]

Support shardy in orbax.

copybara-service[bot]