Revision 473d1c3685570e505bed58512afa05fb7d7a8935 authored by Skye Wanderman-Milne on 28 March 2023, 20:42:51 UTC, committed by Skye Wanderman-Milne on 28 March 2023, 22:28:13 UTC
I forgot that the default setting is actually in jaxlib: https://github.com/openxla/xla/blob/fbe9a80fdb8c429e8a175962459da348cd560a50/xla/python/xla_client.py#L135 To be able to make this change as a jax-only release, I manually set the env var on Cloud TPU if it isn't already set.
1 parent 4061bbb
monitoring.py
# Copyright 2022 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Utilities for instrumenting code.
Code points can be marked as a named event. Every time an event is reached
during program execution, the registered listeners will be invoked.
A typical listener callback is to send an event to a metrics collector for
aggregation/exporting.
"""
from jax._src.monitoring import (
record_event as record_event,
record_event_duration_secs as record_event_duration_secs,
register_event_listener as register_event_listener,
register_event_duration_secs_listener as register_event_duration_secs_listener,
)
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...