<commit_before># Copyright 2018 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.

"""Interface and utility functions to XLA.

This module wraps the XLA client(s) and builders to standardize their interfaces
and provide some automatic type mapping logic for converting between Numpy and
XLA. There are also a handful of related casting utilities.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import warnings
from distutils.util import strtobool

from ..config import flags
import numpy as onp  # 'onp' rather than 'np' to distinguish from autograd.numpy

import jaxlib

# Check the jaxlib version before importing anything else from jaxlib.
def _check_jaxlib_version():
  minimum_version = (0, 1, 11)
  if hasattr(jaxlib, '__version__'):
    version = tuple(int(x) for x in jaxlib.__version__.split('.'))
  else:
    version = (0, 1, 9)  # The version before jaxlib.__version__ was added.
  if version < minimum_version:
    msg = 'jaxlib is version {}, but this version of jax requires version {}.'
    raise ValueError(msg.format('.'.join(map(str, version)),
                                '.'.join(map(str, minimum_version))))

_check_jaxlib_version()


from jaxlib import xla_client
from jaxlib import xla_data_pb2
from jaxlib import xrt<commit_msg>Support older jaxlibs without XRT.<commit_after># Copyright 2018 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.

"""Interface and utility functions to XLA.

This module wraps the XLA client(s) and builders to standardize their interfaces
and provide some automatic type mapping logic for converting between Numpy and
XLA. There are also a handful of related casting utilities.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import warnings
from distutils.util import strtobool

from ..config import flags
import numpy as onp  # 'onp' rather than 'np' to distinguish from autograd.numpy

import jaxlib

# Check the jaxlib version before importing anything else from jaxlib.
def _check_jaxlib_version():
  minimum_version = (0, 1, 11)
  if hasattr(jaxlib, '__version__'):
    version = tuple(int(x) for x in jaxlib.__version__.split('.'))
  else:
    version = (0, 1, 9)  # The version before jaxlib.__version__ was added.
  if version < minimum_version:
    msg = 'jaxlib is version {}, but this version of jax requires version {}.'
    raise ValueError(msg.format('.'.join(map(str, version)),
                                '.'.join(map(str, minimum_version))))

_check_jaxlib_version()


from jaxlib import xla_client
from jaxlib import xla_data_pb2

# Workaround for older jaxlib versions. Remove after a jaxlib release.
try:
  from jaxlib import xrt
except ImportError:
  xrt = None
{'fraction_non_alphanumeric': 0.07204610951008646, 'fraction_numerical': 0.008383547288446425, 'mean_word_length': 4.173441734417344, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5409497', 'n_tokens_mistral': 1078, 'n_tokens_neox': 1044, 'n_words': 476}