def test_create_xla_tff_computation_with_reordered_tensor_indexes(self):
builder = xla_client.XlaBuilder('comp')
tensor_shape_1 = xla_client.Shape.array_shape(
xla_client.dtype_to_etype(np.int32), (10, 1))
param_1 = xla_client.ops.Parameter(builder, 0, tensor_shape_1)
tensor_shape_2 = xla_client.Shape.array_shape(
xla_client.dtype_to_etype(np.int32), (1, 20))
param_2 = xla_client.ops.Parameter(builder, 1, tensor_shape_2)
xla_client.ops.Dot(param_1, param_2)
xla_comp = builder.build()
comp_pb_1 = xla_serialization.create_xla_tff_computation(
xla_comp, [0, 1],
computation_types.FunctionType(
((np.int32, (10, 1)), (np.int32, (1, 20))), (np.int32, (
10,
20,
))))
self.assertIsInstance(comp_pb_1, pb.Computation)
self.assertEqual(comp_pb_1.WhichOneof('computation'), 'xla')
type_spec_1 = type_serialization.deserialize_type(comp_pb_1.type)
self.assertEqual(
str(type_spec_1), '(<int32[10,1],int32[1,20]> -> int32[10,20])')
comp_pb_2 = xla_serialization.create_xla_tff_computation(
xla_comp, [1, 0],
computation_types.FunctionType(
((np.int32, (1, 20)), (np.int32, (10, 1))), (np.int32, (
10,
20,
))))
self.assertIsInstance(comp_pb_2, pb.Computation)
self.assertEqual(comp_pb_2.WhichOneof('computation'), 'xla')
type_spec_2 = type_serialization.deserialize_type(comp_pb_2.type)
self.assertEqual(
str(type_spec_2), '(<int32[1,20],int32[10,1]> -> int32[10,20])')
def test_add_numbers(self):
builder = xla_client.XlaBuilder('comp')
param = xla_client.ops.Parameter(
builder, 0,
xla_client.shape_from_pyval(
tuple([np.array(0, dtype=np.int32)] * 2)))
xla_client.ops.Add(xla_client.ops.GetTupleElement(param, 0),
xla_client.ops.GetTupleElement(param, 1))
xla_comp = builder.build()
comp_type = computation_types.FunctionType((np.int32, np.int32),
np.int32)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_comp, [0, 1], comp_type)
ex = executor.XlaExecutor()
async def _compute_fn():
comp_val = await ex.create_value(comp_pb, comp_type)
x_val = await ex.create_value(20, np.int32)
y_val = await ex.create_value(30, np.int32)
arg_val = await ex.create_struct([x_val, y_val])
call_val = await ex.create_call(comp_val, arg_val)
return await call_val.compute()
result = asyncio.run(_compute_fn())
self.assertEqual(result, 50)
def create_scalar_multiply_operator(
self, operand_type: computation_types.Type,
scalar_type: computation_types.TensorType
) -> local_computation_factory_base.ComputationProtoAndType:
py_typecheck.check_type(operand_type, computation_types.Type)
py_typecheck.check_type(scalar_type, computation_types.TensorType)
if not type_analysis.is_structure_of_tensors(operand_type):
raise ValueError(
'Not a tensor or a structure of tensors: {}'.format(
str(operand_type)))
operand_shapes = _xla_tensor_shape_list_from_from_tff_tensor_or_struct_type(
operand_type)
scalar_shape = _xla_tensor_shape_from_tff_tensor_type(scalar_type)
num_operand_tensors = len(operand_shapes)
builder = xla_client.XlaBuilder('comp')
param = xla_client.ops.Parameter(
builder, 0,
xla_client.Shape.tuple_shape(operand_shapes + [scalar_shape]))
scalar_ref = xla_client.ops.GetTupleElement(param, num_operand_tensors)
result_tensors = []
for idx in range(num_operand_tensors):
result_tensors.append(
xla_client.ops.Mul(xla_client.ops.GetTupleElement(param, idx),
scalar_ref))
xla_client.ops.Tuple(builder, result_tensors)
xla_computation = builder.build()
comp_type = computation_types.FunctionType(
computation_types.StructType([(None, operand_type),
(None, scalar_type)]), operand_type)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_computation, list(range(num_operand_tensors + 1)), comp_type)
return (comp_pb, comp_type)
def _make_test_xla_comp_int32x10_to_int32x10():
builder = xla_client.XlaBuilder('comp')
tensor_shape = xla_client.Shape.array_shape(
xla_client.dtype_to_etype(np.int32), (10,))
param = xla_client.ops.Parameter(builder, 0, tensor_shape)
constant = xla_client.ops.Constant(builder, np.zeros((10,), dtype=np.int32))
xla_client.ops.Add(param, constant)
return builder.build()
def test_xla_computation_and_bindings_to_tff_type_none_binding_to_int32(self):
builder = xla_client.XlaBuilder('comp')
xla_client.ops.Constant(builder, np.int32(10))
xla_computation = builder.build()
parameter_binding = None
result_binding = pb.Xla.Binding(tensor=pb.Xla.TensorBinding(index=0))
tff_type = xla_serialization.xla_computation_and_bindings_to_tff_type(
xla_computation, parameter_binding, result_binding)
self.assertEqual(str(tff_type), '( -> int32)')
def test_set_local_python_execution_context_and_run_simple_xla_computation(
self):
builder = xla_client.XlaBuilder('comp')
xla_client.ops.Parameter(builder, 0,
xla_client.shape_from_pyval(tuple()))
xla_client.ops.Constant(builder, np.int32(10))
xla_comp = builder.build()
comp_type = computation_types.FunctionType(None, np.int32)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_comp, [], comp_type)
ctx_stack = context_stack_impl.context_stack
comp = computation_impl.ConcreteComputation(comp_pb, ctx_stack)
execution_contexts.set_local_python_execution_context()
self.assertEqual(comp(), 10)
def test_to_representation_for_type_with_noarg_to_int32_comp(self):
builder = xla_client.XlaBuilder('comp')
xla_client.ops.Parameter(builder, 0,
xla_client.shape_from_pyval(tuple()))
xla_client.ops.Constant(builder, np.int32(10))
xla_comp = builder.build()
comp_type = computation_types.FunctionType(None, np.int32)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_comp, [], comp_type)
rep = executor.to_representation_for_type(comp_pb, comp_type,
self._backend)
self.assertTrue(callable(rep))
result = rep()
self.assertEqual(result, 10)
def test_computation_callable_return_one_number(self):
builder = xla_client.XlaBuilder('comp')
xla_client.ops.Parameter(builder, 0, xla_client.shape_from_pyval(tuple()))
xla_client.ops.Constant(builder, np.int32(10))
xla_comp = builder.build()
comp_type = computation_types.FunctionType(None, np.int32)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_comp, [], comp_type)
backend = jax.lib.xla_bridge.get_backend()
comp_callable = runtime.ComputationCallable(comp_pb, comp_type, backend)
self.assertIsInstance(comp_callable, runtime.ComputationCallable)
self.assertEqual(str(comp_callable.type_signature), '( -> int32)')
result = comp_callable()
self.assertEqual(result, 10)
def test_xla_computation_and_bindings_to_tff_type_raises_unused_element(self):
builder = xla_client.XlaBuilder('comp')
tensor_shape = xla_client.Shape.array_shape(
xla_client.dtype_to_etype(np.int32), (10,))
tuple_shape = xla_client.Shape.tuple_shape([tensor_shape, tensor_shape])
param = xla_client.ops.Parameter(builder, 0, tuple_shape)
constant = xla_client.ops.Constant(builder, np.zeros((10,), dtype=np.int32))
xla_client.ops.Add(xla_client.ops.GetTupleElement(param, 0), constant)
xla_computation = builder.build()
parameter_binding = pb.Xla.Binding(
struct=pb.Xla.StructBinding(
element=[pb.Xla.Binding(tensor=pb.Xla.TensorBinding(index=0))]))
result_binding = pb.Xla.Binding(tensor=pb.Xla.TensorBinding(index=0))
with self.assertRaises(ValueError):
xla_serialization.xla_computation_and_bindings_to_tff_type(
xla_computation, parameter_binding, result_binding)
def test_xla_computation_and_bindings_to_tff_type_int32_tuple_to_int32(self):
builder = xla_client.XlaBuilder('comp')
tensor_shape = xla_client.Shape.array_shape(
xla_client.dtype_to_etype(np.int32), (10,))
tuple_shape = xla_client.Shape.tuple_shape([tensor_shape])
param = xla_client.ops.Parameter(builder, 0, tuple_shape)
constant = xla_client.ops.Constant(builder, np.zeros((10,), dtype=np.int32))
xla_client.ops.Add(xla_client.ops.GetTupleElement(param, 0), constant)
xla_computation = builder.build()
parameter_binding = pb.Xla.Binding(
struct=pb.Xla.StructBinding(
element=[pb.Xla.Binding(tensor=pb.Xla.TensorBinding(index=0))]))
result_binding = pb.Xla.Binding(tensor=pb.Xla.TensorBinding(index=0))
tff_type = xla_serialization.xla_computation_and_bindings_to_tff_type(
xla_computation, parameter_binding, result_binding)
self.assertEqual(str(tff_type), '(<int32[10]> -> int32[10])')
def test_to_representation_for_type_with_2xint32_to_int32_comp(self):
builder = xla_client.XlaBuilder('comp')
param = xla_client.ops.Parameter(
builder, 0,
xla_client.shape_from_pyval(
tuple([np.array(0, dtype=np.int32)] * 2)))
xla_client.ops.Add(xla_client.ops.GetTupleElement(param, 0),
xla_client.ops.GetTupleElement(param, 1))
xla_comp = builder.build()
comp_type = computation_types.FunctionType((np.int32, np.int32),
np.int32)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_comp, [0, 1], comp_type)
rep = executor.to_representation_for_type(comp_pb, comp_type,
self._backend)
self.assertTrue(callable(rep))
result = rep(
structure.Struct([(None, np.int32(20)), (None, np.int32(30))]))
self.assertEqual(result, 50)
def test_to_representation_for_type_with_noarg_to_2xint32_comp(self):
builder = xla_client.XlaBuilder('comp')
xla_client.ops.Parameter(builder, 0,
xla_client.shape_from_pyval(tuple()))
xla_client.ops.Tuple(builder, [
xla_client.ops.Constant(builder, np.int32(10)),
xla_client.ops.Constant(builder, np.int32(20))
])
xla_comp = builder.build()
comp_type = computation_types.FunctionType(
None,
computation_types.StructType([('a', np.int32), ('b', np.int32)]))
comp_pb = xla_serialization.create_xla_tff_computation(
xla_comp, [0, 1], comp_type)
rep = executor.to_representation_for_type(comp_pb, comp_type,
self._backend)
self.assertTrue(callable(rep))
result = rep()
self.assertEqual(str(result), '<a=10,b=20>')
def _create_xla_binary_op_computation(type_spec, xla_binary_op_constructor):
"""Helper for constructing computations that implement binary operators.
The constructed computation is of type `(<T,T> -> T)`, where `T` is the type
of the operand (`type_spec`).
Args:
type_spec: The type of a single operand.
xla_binary_op_constructor: A two-argument callable that constructs a binary
xla op from tensor parameters (such as `xla_client.ops.Add` or similar).
Returns:
An instance of `local_computation_factory_base.ComputationProtoAndType`.
Raises:
ValueError: if the arguments are invalid.
"""
py_typecheck.check_type(type_spec, computation_types.Type)
if not type_analysis.is_structure_of_tensors(type_spec):
raise ValueError('Not a tensor or a structure of tensors: {}'.format(
str(type_spec)))
tensor_shapes = _xla_tensor_shape_list_from_from_tff_tensor_or_struct_type(
type_spec)
num_tensors = len(tensor_shapes)
builder = xla_client.XlaBuilder('comp')
param = xla_client.ops.Parameter(
builder, 0, xla_client.Shape.tuple_shape(tensor_shapes * 2))
result_tensors = []
for idx in range(num_tensors):
result_tensors.append(
xla_binary_op_constructor(
xla_client.ops.GetTupleElement(param, idx),
xla_client.ops.GetTupleElement(param, idx + num_tensors)))
xla_client.ops.Tuple(builder, result_tensors)
xla_computation = builder.build()
comp_type = computation_types.FunctionType(
computation_types.StructType([(None, type_spec)] * 2), type_spec)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_computation, list(range(2 * num_tensors)), comp_type)
return (comp_pb, comp_type)
def test_computation_callable_add_two_numbers(self):
builder = xla_client.XlaBuilder('comp')
param = xla_client.ops.Parameter(
builder, 0,
xla_client.shape_from_pyval(tuple([np.array(0, dtype=np.int32)] * 2)))
xla_client.ops.Add(
xla_client.ops.GetTupleElement(param, 0),
xla_client.ops.GetTupleElement(param, 1))
xla_comp = builder.build()
comp_type = computation_types.FunctionType((np.int32, np.int32), np.int32)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_comp, [0, 1], comp_type)
backend = jax.lib.xla_bridge.get_backend()
comp_callable = runtime.ComputationCallable(comp_pb, comp_type, backend)
self.assertIsInstance(comp_callable, runtime.ComputationCallable)
self.assertEqual(
str(comp_callable.type_signature), '(<int32,int32> -> int32)')
result = comp_callable(
structure.Struct([(None, np.int32(2)), (None, np.int32(3))]))
self.assertEqual(result, 5)
def create_constant_from_scalar(
self, value, type_spec: computation_types.Type
) -> local_computation_factory_base.ComputationProtoAndType:
py_typecheck.check_type(type_spec, computation_types.Type)
if not type_analysis.is_structure_of_tensors(type_spec):
raise ValueError(
'Not a tensor or a structure of tensors: {}'.format(
str(type_spec)))
builder = xla_client.XlaBuilder('comp')
# We maintain the convention that arguments are supplied as a tuple for the
# sake of consistency and uniformity (see comments in `computation.proto`).
# Since there are no arguments here, we create an empty tuple.
xla_client.ops.Parameter(builder, 0,
xla_client.shape_from_pyval(tuple()))
def _constant_from_tensor(tensor_type):
py_typecheck.check_type(tensor_type, computation_types.TensorType)
numpy_value = np.full(shape=tensor_type.shape.dims,
fill_value=value,
dtype=tensor_type.dtype.as_numpy_dtype)
return xla_client.ops.Constant(builder, numpy_value)
if isinstance(type_spec, computation_types.TensorType):
tensors = [_constant_from_tensor(type_spec)]
else:
tensors = [
_constant_from_tensor(x) for x in structure.flatten(type_spec)
]
# Likewise, results are always returned as a single tuple with results.
# This is always a flat tuple; the nested TFF structure is defined by the
# binding.
xla_client.ops.Tuple(builder, tensors)
xla_computation = builder.build()
comp_type = computation_types.FunctionType(None, type_spec)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_computation, [], comp_type)
return (comp_pb, comp_type)
def test_create_and_invoke_noarg_comp_returning_int32(self):
builder = xla_client.XlaBuilder('comp')
xla_client.ops.Parameter(builder, 0,
xla_client.shape_from_pyval(tuple()))
xla_client.ops.Constant(builder, np.int32(10))
xla_comp = builder.build()
comp_type = computation_types.FunctionType(None, np.int32)
comp_pb = xla_serialization.create_xla_tff_computation(
xla_comp, [], comp_type)
ex = executor.XlaExecutor()
comp_val = asyncio.run(ex.create_value(comp_pb, comp_type))
self.assertIsInstance(comp_val, executor.XlaValue)
self.assertEqual(str(comp_val.type_signature), str(comp_type))
self.assertTrue(callable(comp_val.internal_representation))
result = comp_val.internal_representation()
self.assertEqual(result, 10)
call_val = asyncio.run(ex.create_call(comp_val))
self.assertIsInstance(call_val, executor.XlaValue)
self.assertEqual(str(call_val.type_signature), 'int32')
result = asyncio.run(call_val.compute())
self.assertEqual(result, 10)
def _make_test_xla_comp_noarg_to_int32():
builder = xla_client.XlaBuilder('comp')
xla_client.ops.Constant(builder, np.int32(10))
return builder.build()