<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.

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

from collections import namedtuple
import itertools as it
import numpy as onp
import operator as op
import six
from six.moves import xrange

from ..config import flags
from .. import core
from .. import ad_util
from ..abstract_arrays import ConcreteArray, ShapedArray, make_shaped_array, array_types
from ..core import AbstractTuple, JaxTuple, pack, valid_jaxtype
from ..util import partial, partialmethod, memoize, unzip2, concatenate, safe_map, prod
from ..linear_util import transformation_with_aux, memoize as linear_memoize
from ..lib import xla_bridge as xb
from .partial_eval import trace_to_subjaxpr, merge_pvals, JaxprTrace, PartialVal

FLAGS = flags.FLAGS
flags.DEFINE_bool('jax_device_values', True, 'Enable device-persistent values.')

map = safe_map

def apply_primitive(prim, *args, **kwargs):
  abstract_args = map(abstractify, args)
  compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
  return compiled_fun(*args)

@memoize
def xla_primitive_callable(prim, *abstract_args, **kwargs):
  shapes = map(xla_shape, abstract_args)
  built_c = primitive_computation(prim, *shapes, **kwargs)
  result_shape = xla_shape_to_result_shape(built_c.GetReturnValueShape())
  handle_result = result_handler(result_shape)
  compiled = built_c.Compile(shapes, xb.get_compile_options())
  return partial(execute_compiled_primitive, compiled, handle_result)

@memoize
def primitive_computation(prim, *shapes, **kwargs):
  c = xb.make_computation_builder("primitive_computation")
  xla_args = map(c.ParameterWithShape, shapes)
  xla_result = translation_rule(prim)(c, *xla_args, **kwargs)
  try:
    return c.Build()
  except RuntimeError as e:
    # try for a better error message by using the abstract_eval checks
    prim.abstract_eval(*map(aval_from_xla_shape, shapes), **kwargs)
    raise e

def aval_from_xla_shape(shape):
  return ShapedArray(shape.dimensions(), shape.element_type())

def execute_compiled_primitive(compiled, result_handler, *args):
  input_bufs = [device_put(canonicalize_pyval_dtype(x)) for x in args]
  return result_handler(compiled.Execute(input_bufs, not core.skip_checks))

def device_put(x):
  if type(x) is DeviceArray:
    return x.device_buffer
  else:
    return xb.device_put(x)  # can round-trip elements of tuples here

def result_handler(result_shape):
  if type(result_shape) is ResultArray and FLAGS.jax_device_values:
    def handle_result(device_buffer):
      return DeviceArray(device_buffer, *result_shape)
  elif type(result_shape) is ResultArray:
    def handle_result(device_buffer):
      return onp.asarray(DeviceArray(device_buffer, *result_shape))
  elif type(result_shape) is ResultTuple:
    handlers = list(map(result_handler, result_shape))
    def handle_result(device_buffer):
      bufs = device_buffer.destructure()
      return JaxTuple(handler(buf) for handler, buf in zip(handlers, bufs))
  else:
    raise TypeError(type(result_shape))
  return handle_result

def xla_shape_to_result_shape(xla_shape):
  if xla_shape.is_tuple():
    return ResultTuple(map(xla_shape_to_result_shape, xla_shape.tuple_shapes()))
  else:
    shape, dtype = xla_shape.dimensions(), xla_shape.element_type()
    ndim, size = len(shape), prod(shape)
    return ResultArray((shape, dtype, ndim, size))
class ResultTuple(tuple): pass
class ResultArray(tuple): pass


def compile_jaxpr(jaxpr, const_vals, *abstract_args):
  arg_shapes = list(map(xla_shape, abstract_args))
  built_c = jaxpr_computation(jaxpr, const_vals, (), *arg_shapes)
  result_shape = xla_shape_to_result_shape(built_c.GetReturnValueShape())
  return built_c.Compile(arg_shapes, xb.get_compile_options()), result_shape

def jaxpr_computation(jaxpr, const_vals, freevar_shapes, *arg_shapes):
  c = xb.make_computation_builder("jaxpr_computation")

  def read(v):
    return env[v]

  def write(v, node):
    assert node is not None
    env[v] = node

  env = {}
  consts_env = dict(zip(jaxpr.constvars, const_vals))
  write(core.unitvar, c.Tuple())
  map(write, jaxpr.constvars, map(c.Constant, const_vals))
  map(write, jaxpr.freevars, map(c.ParameterWithShape, freevar_shapes))
  map(write, jaxpr.invars, map(c.ParameterWithShape, arg_shapes))
  for eqn in jaxpr.eqns:
    in_nodes = map(read, eqn.invars)
    in_shapes = map(c.GetShape, in_nodes)
    subcs = [jaxpr_computation(subjaxpr,
                               [consts_env[b] for b in const_bindings],
                               map(c.GetShape, map(read, freevar_bindings)),
                               *in_shapes)
             for subjaxpr, const_bindings, freevar_bindings
             in eqn.bound_subjaxprs]
    subfuns = [(subc, tuple(map(read, freevar_bindings)))
               for subc, (_, _, freevar_bindings)
               in zip(subcs, eqn.bound_subjaxprs)]
    ans = translation_rule(eqn.primitive)(c, *(subfuns + in_nodes), **eqn.params)
    out_nodes = xla_destructure(c, ans) if eqn.destructure else [ans]
    map(write, eqn.outvars, out_nodes)
  return c.Build(read(jaxpr.outvar))

def xla_destructure(c, ans):
  num_elements = len(c.GetShape(ans).tuple_shapes())
  return [c.GetTupleElement(ans, i) for i in range(num_elements)]

def unit_constant(c, val):
  assert not val  # must be unit
  return c.Tuple()
xb.register_constant_handler(JaxTuple, unit_constant)

def translation_rule(p):
  try:
    return translations[p]
  except KeyError:
    raise NotImplementedError(
        "XLA translation rule for '{}' not implemented".format(p))


translations = {}

translations[core.pack_p] = lambda c, *xs: c.Tuple(*xs)
translations[ad_util.add_jaxvals_p] = lambda c, x, y: c.Add(x, y)
translations[core.call_p] = lambda c, subc_a1, *a2: c.Call(subc_a1[0],
                                                           subc_a1[1] + a2)
translations[core.identity_p] = lambda c, x: x


def canonicalize_pyval_dtype(x):
  try:
    return canonicalize_dtype_handlers[type(x)](x)<commit_msg>add ad_util.zeros_like translation rule PAIR=hawkinsp<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.

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

from collections import namedtuple
import itertools as it
import numpy as onp
import operator as op
import six
from six.moves import xrange

from ..config import flags
from .. import core
from .. import ad_util
from ..abstract_arrays import ConcreteArray, ShapedArray, make_shaped_array, array_types
from ..core import AbstractTuple, JaxTuple, pack, valid_jaxtype
from ..util import partial, partialmethod, memoize, unzip2, concatenate, safe_map, prod
from ..linear_util import transformation_with_aux, memoize as linear_memoize
from ..lib import xla_bridge as xb
from .partial_eval import trace_to_subjaxpr, merge_pvals, JaxprTrace, PartialVal

FLAGS = flags.FLAGS
flags.DEFINE_bool('jax_device_values', True, 'Enable device-persistent values.')

map = safe_map

def apply_primitive(prim, *args, **kwargs):
  abstract_args = map(abstractify, args)
  compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
  return compiled_fun(*args)

@memoize
def xla_primitive_callable(prim, *abstract_args, **kwargs):
  shapes = map(xla_shape, abstract_args)
  built_c = primitive_computation(prim, *shapes, **kwargs)
  result_shape = xla_shape_to_result_shape(built_c.GetReturnValueShape())
  handle_result = result_handler(result_shape)
  compiled = built_c.Compile(shapes, xb.get_compile_options())
  return partial(execute_compiled_primitive, compiled, handle_result)

@memoize
def primitive_computation(prim, *shapes, **kwargs):
  c = xb.make_computation_builder("primitive_computation")
  xla_args = map(c.ParameterWithShape, shapes)
  xla_result = translation_rule(prim)(c, *xla_args, **kwargs)
  try:
    return c.Build()
  except RuntimeError as e:
    # try for a better error message by using the abstract_eval checks
    prim.abstract_eval(*map(aval_from_xla_shape, shapes), **kwargs)
    raise e

def aval_from_xla_shape(shape):
  return ShapedArray(shape.dimensions(), shape.element_type())

def execute_compiled_primitive(compiled, result_handler, *args):
  input_bufs = [device_put(canonicalize_pyval_dtype(x)) for x in args]
  return result_handler(compiled.Execute(input_bufs, not core.skip_checks))

def device_put(x):
  if type(x) is DeviceArray:
    return x.device_buffer
  else:
    return xb.device_put(x)  # can round-trip elements of tuples here

def result_handler(result_shape):
  if type(result_shape) is ResultArray and FLAGS.jax_device_values:
    def handle_result(device_buffer):
      return DeviceArray(device_buffer, *result_shape)
  elif type(result_shape) is ResultArray:
    def handle_result(device_buffer):
      return onp.asarray(DeviceArray(device_buffer, *result_shape))
  elif type(result_shape) is ResultTuple:
    handlers = list(map(result_handler, result_shape))
    def handle_result(device_buffer):
      bufs = device_buffer.destructure()
      return JaxTuple(handler(buf) for handler, buf in zip(handlers, bufs))
  else:
    raise TypeError(type(result_shape))
  return handle_result

def xla_shape_to_result_shape(xla_shape):
  if xla_shape.is_tuple():
    return ResultTuple(map(xla_shape_to_result_shape, xla_shape.tuple_shapes()))
  else:
    shape, dtype = xla_shape.dimensions(), xla_shape.element_type()
    ndim, size = len(shape), prod(shape)
    return ResultArray((shape, dtype, ndim, size))
class ResultTuple(tuple): pass
class ResultArray(tuple): pass


def compile_jaxpr(jaxpr, const_vals, *abstract_args):
  arg_shapes = list(map(xla_shape, abstract_args))
  built_c = jaxpr_computation(jaxpr, const_vals, (), *arg_shapes)
  result_shape = xla_shape_to_result_shape(built_c.GetReturnValueShape())
  return built_c.Compile(arg_shapes, xb.get_compile_options()), result_shape

def jaxpr_computation(jaxpr, const_vals, freevar_shapes, *arg_shapes):
  c = xb.make_computation_builder("jaxpr_computation")

  def read(v):
    return env[v]

  def write(v, node):
    assert node is not None
    env[v] = node

  env = {}
  consts_env = dict(zip(jaxpr.constvars, const_vals))
  write(core.unitvar, c.Tuple())
  map(write, jaxpr.constvars, map(c.Constant, const_vals))
  map(write, jaxpr.freevars, map(c.ParameterWithShape, freevar_shapes))
  map(write, jaxpr.invars, map(c.ParameterWithShape, arg_shapes))
  for eqn in jaxpr.eqns:
    in_nodes = map(read, eqn.invars)
    in_shapes = map(c.GetShape, in_nodes)
    subcs = [jaxpr_computation(subjaxpr,
                               [consts_env[b] for b in const_bindings],
                               map(c.GetShape, map(read, freevar_bindings)),
                               *in_shapes)
             for subjaxpr, const_bindings, freevar_bindings
             in eqn.bound_subjaxprs]
    subfuns = [(subc, tuple(map(read, freevar_bindings)))
               for subc, (_, _, freevar_bindings)
               in zip(subcs, eqn.bound_subjaxprs)]
    ans = translation_rule(eqn.primitive)(c, *(subfuns + in_nodes), **eqn.params)
    out_nodes = xla_destructure(c, ans) if eqn.destructure else [ans]
    map(write, eqn.outvars, out_nodes)
  return c.Build(read(jaxpr.outvar))

def xla_destructure(c, ans):
  num_elements = len(c.GetShape(ans).tuple_shapes())
  return [c.GetTupleElement(ans, i) for i in range(num_elements)]

def unit_constant(c, val):
  assert not val  # must be unit
  return c.Tuple()
xb.register_constant_handler(JaxTuple, unit_constant)

def translation_rule(p):
  try:
    return translations[p]
  except KeyError:
    raise NotImplementedError(
        "XLA translation rule for '{}' not implemented".format(p))


translations = {}

translations[core.pack_p] = lambda c, *xs: c.Tuple(*xs)
translations[core.call_p] = lambda c, subc_a1, *a2: c.Call(subc_a1[0],
                                                           subc_a1[1] + a2)
translations[core.identity_p] = lambda c, x: x

# TODO(mattjj): zeros_like and add_jaxvals should handle any jaxval
def zeros_like_translation_rule(c, x):
  x_shape = c.GetShape(x)
  return c.Broadcast(c.Constant(onp.array(0, x_shape.element_type())),
                     x_shape.dimensions())
translations[ad_util.zeros_like_p] = zeros_like_translation_rule
translations[ad_util.add_jaxvals_p] = lambda c, x, y: c.Add(x, y)


def canonicalize_pyval_dtype(x):
  try:
    return canonicalize_dtype_handlers[type(x)](x)
{'fraction_non_alphanumeric': 0.097984198478919, 'fraction_numerical': 0.002436683157350661, 'mean_word_length': 4.608281573498965, 'pattern_counts': {'":': 0, '<': 3, '<?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': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2063435', 'n_tokens_mistral': 4500, 'n_tokens_neox': 4435, 'n_words': 1108}