CommonLoopUtils

CommonLoopUtils copied to clipboard

Update clu.Metrics.compute_value() return type, to avoid signature-mismatch error.

1

Update clu.Metrics.compute_value() return type, to avoid signature-mismatch error. according to here: go/pax/metrics#compute-value-return-types

copybara-service[bot]

This PR makes `Metric` and `Collection` inherit from `flax.struct.PyTreeNode`, this makes creating new metric types more ergonomic as you no longer need to decorate the class with `flax.struct.dataclass`.

cgarciae

Warmup input pipeline in the background.

copybara-service[bot]

Add point clouds summary writer to tensorboard interface and metric writer.

1

Add point clouds summary writer to tensorboard interface and metric writer.

copybara-service[bot]

Make g3pdb lazy again

copybara-service[bot]

Make sklearn import lazy

copybara-service[bot]

Rename differentially_private_aggregate -> optax.contrib

copybara-service[bot]

Adds `reshuffle_each_iteration` argument to `deterministic_data.create_dataset()`. This argument is passed to `tf.data.Dataset.shuffle()` and controls whether the dataset is reshuffled each time it is iterated over. The default value is `None`, which...

copybara-service[bot]

Silence some pytype errors related to a JAX build refactor This build change allows pytype to propagate annotations that it previously did not, and because of this it starts flagging...

copybara-service[bot]

Update references to JAX's GitHub repo JAX has moved from https://github.com/google/jax to https://github.com/jax-ml/jax

copybara-service[bot]