# JAX dtypes differ from NumPy in both: # a) their type promotion rules, and # b) the set of supported types (e.g., bfloat16), # so we need our own implementation that deviates from NumPy in places. import functools from typing import Any, Dict import numpy as np from jax._src import util from jax._src.config import flags, config from jax.lib import xla_client from jax._src import traceback_util traceback_util.register_exclusion(__file__) FLAGS = flags.FLAGS # bfloat16 support bfloat16: type = xla_client.bfloat16 _bfloat16_dtype: np.dtype = np.dtype(bfloat16) # Default types. bool_ = np.bool_ int_: np.dtype = np.int64 # type: ignore float_: np.dtype = np.float64 # type: ignore complex_ = np.complex128 # TODO(phawkins): change the above defaults to:
# Copyright 2020 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. import os from jax._src import traceback_util traceback_util.register_exclusion(os.path.dirname(__file__))