Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

metrics: Use Array and ArrayLike types thoughout #278

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

copybara-service[bot]
Copy link

@copybara-service copybara-service bot commented May 4, 2023

metrics: Use Array and ArrayLike types thoughout

Currently the inputs to from_model_output are not typed. However,
these functions cannot accept arbitrary inputs, they need to be a value
convertable to a jax.Array. This change fixes this so that:

  • from_model_output takes in types of Array or ArrayLike
  • Removes use of jnp.array as a type as it's equivalent to Any
  • Makes members of Metric classes have type Array
  • Moves mask checking code into its own function

While we could make everything use Array (instead of ArrayLike),
this would break code like:

@flax.struct.dataclass
class Collection(metrics.Collection):
  train_accuracy: metrics.Accuracy
  learning_rate: metrics.LastValue.from_output("learning_rate")

Collection.gather_from_model_output(learning_rate=0.02, ...)

which seems undesirable.

Note that count and value for LastValue have type ArrayLike,
as this code needs to support passing a plain number for value or
count. Also, the base Metric.compute() method has type Any,
because some metrics return Array while others use dict[str, Array].

Currently the inputs to `from_model_output` are not typed. However,
these functions cannot accept arbitrary inputs, they need to be a value
convertable to a `jax.Array`. This change fixes this so that:
  - `from_model_output` takes in types of `Array` or `ArrayLike`
  - Removes use of `jnp.array` as a type as it's equivalent to `Any`
  - Makes members of Metric classes have type `Array`
  - Moves mask checking code into its own function

While we could make everything use `Array` (instead of `ArrayLike`),
this would break code like:
```
@flax.struct.dataclass
class Collection(metrics.Collection):
  train_accuracy: metrics.Accuracy
  learning_rate: metrics.LastValue.from_output("learning_rate")

Collection.gather_from_model_output(learning_rate=0.02, ...)
```
which seems undesirable.

Note that `count` and `value` for `LastValue` have type `ArrayLike`,
as this code needs to support passing a plain number for `value` or
`count`. Also, the base `Metric.compute()` method has type `Any`,
because some metrics return `Array` while others use `dict[str, Array]`.

PiperOrigin-RevId: 529227218
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant