mirror of
https://github.com/ROCm/jax.git
synced 2025-04-24 06:36:05 +00:00
60e0e9f929
Introduce a config flag for upgrading to a world of custom PRNGs. The flag defaults off, so that we can introduce custom PRNGs into the codebase and allow downstream libraries time to upgrade. Backwards compatible behavior is meant in an external sense. This does not mean that our code is internally the same any longer.
24 lines
715 B
Python
24 lines
715 B
Python
# Copyright 2021 Google LLC
|
|
#
|
|
# 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.
|
|
|
|
# flake8: noqa: F401
|
|
|
|
from jax._src.prng import (
|
|
PRNGImpl,
|
|
seed_with_impl,
|
|
threefry2x32_p,
|
|
threefry_2x32,
|
|
threefry_prng_impl,
|
|
)
|