EVOLUTION-MANAGER

Edit File: xla.py

# Copyright 2018 The TensorFlow Authors. All Rights Reserved. # # 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 # # http://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. # ============================================================================== """Experimental library that exposes XLA operations directly in TensorFlow. It is sometimes useful to be able to build HLO programs directly from TensorFlow. This file provides Tensorflow operators that mirror the semantics of HLO operators as closely as possible. Note: There is no promise of backward or forward compatibility for operators defined in this module. This is primarily because the underlying HLO operators do not promise backward or forward compatibility. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function from tensorflow.compiler.tf2xla.ops import gen_xla_ops from tensorflow.core.framework import attr_value_pb2 from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import bitwise_ops from tensorflow.python.ops import gen_math_ops from tensorflow.python.ops import gen_random_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import special_math_ops # TODO(phawkins): provide wrappers for all XLA operators. Currently the missing # ops include: # infeed/outfeed (available via tf.contrib.tpu) # collectives, e.g., cross-replica-sum (available via tf.contrib.tpu) # conditional # gather/scatter # collapse # This file reuses builtin names (following XLA's names, so we can call things # like xla.max), so we capture the builtin versions here. # pylint: disable=redefined-builtin _max = max _min = min _slice = slice # pylint: disable=invalid-name constant = constant_op.constant # Unary operators. # For most arithmetic operators there is a TensorFlow operator # that exactly corresponds to each XLA operator. Rather than defining # XLA-specific variants, we reuse the corresponding TensorFlow operator. # TODO(phawkins): It would be even better to have TensorFlow operators that 1:1 # wrap every HLO operator, because that would allow us to be confident that the # semantics match. def _unary_op(fn): """Wrapper that restricts `fn` to have the correct signature.""" def unary_op_wrapper(x, name=None): return fn(x, name=name) return unary_op_wrapper abs = _unary_op(math_ops.abs) # TODO(phawkins): implement clz. conj = _unary_op(math_ops.conj) cos = _unary_op(math_ops.cos) ceil = _unary_op(math_ops.ceil) digamma = _unary_op(math_ops.digamma) erf = _unary_op(math_ops.erf) erfc = _unary_op(math_ops.erfc) erfinv = _unary_op(math_ops.erfinv) ndtri = _unary_op(math_ops.ndtri) exp = _unary_op(math_ops.exp) expm1 = _unary_op(math_ops.expm1) floor = _unary_op(math_ops.floor) imag = _unary_op(math_ops.imag) is_finite = _unary_op(math_ops.is_finite) lgamma = _unary_op(math_ops.lgamma) log = _unary_op(math_ops.log) log1p = _unary_op(math_ops.log1p) logical_not = _unary_op(math_ops.logical_not) neg = _unary_op(math_ops.neg) real = _unary_op(math_ops.real) # TODO(phawkins): unlike xla::Round, this rounds to even instead of zero for # numbers halfway between two integers. round = _unary_op(math_ops.round) sin = _unary_op(math_ops.sin) sign = _unary_op(math_ops.sign) tanh = _unary_op(math_ops.tanh) # Bessel bessel_i0e = _unary_op(special_math_ops.bessel_i0e) bessel_i1e = _unary_op(special_math_ops.bessel_i1e) # Binary operators # The main difference between TensorFlow and XLA binary ops is the broadcasting # semantics. TensorFlow uses Numpy-style broadcasting semantics, whereas XLA # requires an explicit specification of which dimensions to broadcast if the # arguments have different ranks. def _broadcasting_binary_op(fn): """Wraps a binary Tensorflow operator and performs XLA-style broadcasting.""" def broadcasting_binary_op_wrapper(x, y, broadcast_dims=None, name=None): """Inner wrapper function.""" broadcast_dims = broadcast_dims or [] broadcast_dims = ops.convert_to_tensor(broadcast_dims, dtypes.int64) # Rather than relying on having static shape information in the TensorFlow # graph, we use an XlaBroadcastHelper op that can compute the correct shapes # at JIT compilation time. x, y = gen_xla_ops.xla_broadcast_helper(x, y, broadcast_dims) return fn(x, y, name=name) return broadcasting_binary_op_wrapper # Map from TF signed types to TF unsigned types. _SIGNED_TO_UNSIGNED_TABLE = { dtypes.int8: dtypes.uint8, dtypes.int16: dtypes.uint16, dtypes.int32: dtypes.uint32, dtypes.int64: dtypes.uint64, } # Map from TF unsigned types to TF signed types. _UNSIGNED_TO_SIGNED_TABLE = { dtypes.uint8: dtypes.int8, dtypes.uint16: dtypes.int16, dtypes.uint32: dtypes.int32, dtypes.uint64: dtypes.int64, } def _shift_right_logical_helper(x, y, name=None): """Performs an integer right logical shift irrespective of input type.""" assert y.dtype == x.dtype dtype = x.dtype signed = dtype in _SIGNED_TO_UNSIGNED_TABLE if signed: unsigned_dtype = _SIGNED_TO_UNSIGNED_TABLE[dtype] x = math_ops.cast(x, unsigned_dtype) y = math_ops.cast(y, unsigned_dtype) output = bitwise_ops.right_shift(x, y, name=name) if signed: output = math_ops.cast(output, dtype) return output def _shift_right_arithmetic_helper(x, y, name=None): """Performs an integer right arithmetic shift irrespective of input type.""" assert y.dtype == x.dtype dtype = x.dtype unsigned = dtype in _UNSIGNED_TO_SIGNED_TABLE if unsigned: signed_dtype = _UNSIGNED_TO_SIGNED_TABLE[dtype] x = math_ops.cast(x, signed_dtype) y = math_ops.cast(y, signed_dtype) output = bitwise_ops.right_shift(x, y, name=name) if unsigned: output = math_ops.cast(output, dtype) return output add = _broadcasting_binary_op(math_ops.add) sub = _broadcasting_binary_op(math_ops.sub) mul = _broadcasting_binary_op(math_ops.mul) div = _broadcasting_binary_op(math_ops.div) rem = _broadcasting_binary_op(gen_math_ops.mod) max = _broadcasting_binary_op(math_ops.maximum) min = _broadcasting_binary_op(math_ops.minimum) atan2 = _broadcasting_binary_op(math_ops.atan2) complex = _broadcasting_binary_op(math_ops.complex) logical_and = _broadcasting_binary_op(math_ops.logical_and) logical_or = _broadcasting_binary_op(math_ops.logical_or) logical_xor = _broadcasting_binary_op(math_ops.logical_xor) eq = _broadcasting_binary_op(math_ops.equal) ne = _broadcasting_binary_op(math_ops.not_equal) ge = _broadcasting_binary_op(math_ops.greater_equal) gt = _broadcasting_binary_op(math_ops.greater) le = _broadcasting_binary_op(math_ops.less_equal) lt = _broadcasting_binary_op(math_ops.less) pow = _broadcasting_binary_op(math_ops.pow) shift_left = _broadcasting_binary_op(bitwise_ops.left_shift) shift_right_logical = _broadcasting_binary_op(_shift_right_logical_helper) shift_right_arithmetic = _broadcasting_binary_op(_shift_right_arithmetic_helper) igamma = _broadcasting_binary_op(math_ops.igamma) igamma_grad_a = _broadcasting_binary_op(gen_math_ops.igamma_grad_a) random_gamma_grad = _broadcasting_binary_op(gen_random_ops.random_gamma_grad) igammac = _broadcasting_binary_op(math_ops.igammac) polygamma = _broadcasting_binary_op(math_ops.polygamma) zeta = _broadcasting_binary_op(math_ops.zeta) def _binary_op(fn): """Wrapper that restricts `fn` to have the correct signature.""" def binary_op_wrapper(x, y, name=None): return fn(x, y, name=name) return binary_op_wrapper transpose = _binary_op(array_ops.transpose) rev = _binary_op(array_ops.reverse) bitcast_convert_type = array_ops.bitcast def broadcast(x, dims, name=None): x = ops.convert_to_tensor(x) shape = array_ops.concat([constant_op.constant(dims), array_ops.shape(x)], axis=0) return array_ops.broadcast_to(x, shape, name=name) def clamp(a, x, b, name=None): return min(max(a, x, name=name), b, name=name) concatenate = array_ops.concat def conv(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count=1, precision_config=None, name=None): """Wraps the XLA ConvGeneralDilated operator. ConvGeneralDilated is the most general form of XLA convolution and is documented at https://www.tensorflow.org/performance/xla/operation_semantics#conv_convolution Args: lhs: the input tensor rhs: the kernel tensor window_strides: the inter-window strides padding: the padding to apply at the start and end of each input dimensions lhs_dilation: dilation to apply between input elements rhs_dilation: dilation to apply between kernel elements dimension_numbers: a `ConvolutionDimensionNumbers` proto. feature_group_count: number of feature groups for grouped convolution. precision_config: a `xla.PrecisionConfig` proto. name: an optional name for the operator Returns: A tensor representing the output of the convolution. """ precision_config_proto = "" if precision_config: precision_config_proto = precision_config.SerializeToString() return gen_xla_ops.xla_conv( lhs, rhs, window_strides=window_strides, padding=padding, lhs_dilation=lhs_dilation, rhs_dilation=rhs_dilation, feature_group_count=feature_group_count, dimension_numbers=dimension_numbers.SerializeToString(), precision_config=precision_config_proto, name=name) convert_element_type = math_ops.cast def dot(lhs, rhs, name=None): return math_ops.tensordot(lhs, rhs, axes=1, name=name) def dot_general(lhs, rhs, dimension_numbers, precision_config=None, name=None): precision_config_proto = "" if precision_config: precision_config_proto = precision_config.SerializeToString() return gen_xla_ops.xla_dot( lhs, rhs, dimension_numbers=dimension_numbers.SerializeToString(), precision_config=precision_config_proto, name=name) def self_adjoint_eig(a, lower, max_iter, epsilon): return gen_xla_ops.xla_self_adjoint_eig(a, lower, max_iter, epsilon) def svd(a, max_iter, epsilon, precision_config=None): precision_config_proto = "" if precision_config: precision_config_proto = precision_config.SerializeToString() return gen_xla_ops.xla_svd(a, max_iter, epsilon, precision_config_proto) dynamic_slice = gen_xla_ops.xla_dynamic_slice dynamic_update_slice = gen_xla_ops.xla_dynamic_update_slice einsum = gen_xla_ops.xla_einsum # TODO(phawkins): generalize tf.pad to support interior padding, and then remove # the XLA-specific pad operator. pad = gen_xla_ops.xla_pad def random_normal(mu, sigma, dims, name=None): mu = ops.convert_to_tensor(mu) return random_ops.random_normal( dims, mean=mu, stddev=sigma, dtype=mu.dtype, name=name) def random_uniform(minval, maxval, dims, name=None): minval = ops.convert_to_tensor(minval) return random_ops.random_uniform( dims, minval, maxval, dtype=minval.dtype, name=name) recv = gen_xla_ops.xla_recv reduce = gen_xla_ops.xla_reduce variadic_reduce = gen_xla_ops.xla_variadic_reduce def reduce_window(operand, init, reducer, window_dimensions, window_strides=None, base_dilations=None, window_dilations=None, padding=None, name=None): """Wraps the XLA ReduceWindow operator. ReduceWindow is documented at https://www.tensorflow.org/performance/xla/operation_semantics#reducewindow . Args: operand: the input tensor init: a scalar tensor representing the initial value for the reduction reducer: a reduction function that combines a pair of scalars. window_dimensions: shape of the window, as a list of integers window_strides: inter-window strides, as a list of integers. Optional; if omitted, defaults to strides of 1. padding: padding to apply to 'operand'. List of (low, high) pairs of integers that specify the padding to apply before and after each dimension. Optional; if omitted, defaults to no padding. name: the operator name, or None. Returns: A tensor that represents the output of the reduce_window operator. """ window_strides = window_strides or [1] * len(window_dimensions) base_dilations = base_dilations or [1] * len(window_dimensions) window_dilations = window_dilations or [1] * len(window_dimensions) padding = padding or [(0, 0)] * len(window_dimensions) return gen_xla_ops.xla_reduce_window( input=operand, init_value=init, window_dimensions=window_dimensions, window_strides=window_strides, base_dilations=base_dilations, window_dilations=window_dilations, padding=padding, computation=reducer, name=name) replica_id = gen_xla_ops.xla_replica_id # Set a static bound for the given input value as a hint to Xla compiler, # returns the same value. # Usage: # def f(t, p): # p = xla.set_bound(p, 3) # Tells xla the constraint that p <= 3. # return t[:p] # xla knows the bound of the slice is 3. set_bound = gen_xla_ops.xla_set_bound def reshape(x, new_sizes, dimensions=None, name=None): if dimensions is not None: x = array_ops.transpose(x, dimensions) x = array_ops.reshape(x, new_sizes, name=name) return x def select(condition, x, y, name=None): return array_ops.where(condition, x, y, name) select_and_scatter = gen_xla_ops.xla_select_and_scatter send = gen_xla_ops.xla_send def slice(x, start_dims, limit_dims, strides): spec = [ _slice(start, limit, stride) for (start, limit, stride) in zip(start_dims, limit_dims, strides) ] return x[tuple(spec)] sharding = gen_xla_ops.xla_sharding @ops.RegisterGradient("XlaSharding") def _sharding_grad(op, grad): grad_sharding = gen_xla_ops.xla_sharding(grad) # pylint: disable=protected-access grad_sharding.op._set_attr( "_XlaSharding", attr_value_pb2.AttrValue(s=op.get_attr("_XlaSharding"))) return [grad_sharding] spmd_full_to_shard_shape = gen_xla_ops.xla_spmd_full_to_shard_shape spmd_shard_to_full_shape = gen_xla_ops.xla_spmd_shard_to_full_shape @ops.RegisterGradient("XlaSpmdFullToShardShape") def _spmd_full_to_shard_shape_grad(op, grad): s2f = gen_xla_ops.xla_spmd_shard_to_full_shape( grad, manual_sharding=op.get_attr("manual_sharding"), full_shape=op.inputs[0].shape.as_list()) return [s2f] @ops.RegisterGradient("XlaSpmdShardToFullShape") def _spmd_shard_to_full_shape_grad(op, grad): f2s = gen_xla_ops.xla_spmd_full_to_shard_shape( grad, manual_sharding=op.get_attr("manual_sharding")) return [f2s] sort = gen_xla_ops.xla_sort key_value_sort = gen_xla_ops.xla_key_value_sort while_loop = gen_xla_ops.xla_while dequantize = gen_xla_ops.xla_dequantize def gather(operand, start_indices, dimension_numbers, slice_sizes, indices_are_sorted=False, name=None): return gen_xla_ops.xla_gather( operand, start_indices, slice_sizes=slice_sizes, dimension_numbers=dimension_numbers.SerializeToString(), indices_are_sorted=indices_are_sorted, name=name) def scatter(operand, scatter_indices, updates, update_computation, dimension_numbers, indices_are_sorted=False, name=None): return gen_xla_ops.xla_scatter( operand, scatter_indices, updates, update_computation=update_computation, dimension_numbers=dimension_numbers.SerializeToString(), indices_are_sorted=indices_are_sorted, name=name)