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, 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.get_event_loop().run_until_complete(_compute_fn())
self.assertEqual(result, 50)
def serialize_jax_computation(traced_fn, arg_fn, parameter_type,
context_stack):
"""Serializes a Python function containing JAX code as a TFF computation.
Args:
traced_fn: The Python function containing JAX code to be traced by JAX and
serialized as a TFF computation containing XLA code.
arg_fn: An unpacking function that takes a TFF argument, and returns a combo
of (args, kwargs) to invoke `traced_fn` with (e.g., as the one constructed
by `function_utils.create_argument_unpacking_fn`).
parameter_type: An instance of `computation_types.Type` that represents the
TFF type of the computation parameter, or `None` if the function does not
take any parameters.
context_stack: The context stack to use during serialization.
Returns:
An instance of `pb.Computation` with the constructed computation.
Raises:
TypeError: if the arguments are of the wrong types.
"""
py_typecheck.check_callable(traced_fn)
py_typecheck.check_callable(arg_fn)
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
if parameter_type is not None:
parameter_type = computation_types.to_type(parameter_type)
packed_arg = _tff_type_to_xla_serializer_arg(parameter_type)
else:
packed_arg = None
args, kwargs = arg_fn(packed_arg)
def _adjust_arg(x):
return type_conversions.type_to_py_container(x, x.type_signature)
args = [_adjust_arg(x) for x in args]
kwargs = {k: _adjust_arg(v) for k, v in kwargs.items()}
context = jax_computation_context.JaxComputationContext()
with context_stack.install(context):
tracer_callable = jax.xla_computation(traced_fn,
tuple_args=True,
return_shape=True)
compiled_xla, returned_shape = tracer_callable(*args, **kwargs)
if isinstance(returned_shape, jax.ShapeDtypeStruct):
returned_type_spec = _jax_shape_dtype_struct_to_tff_tensor(
returned_shape)
else:
returned_type_spec = computation_types.to_type(
structure.map_structure(
_jax_shape_dtype_struct_to_tff_tensor,
structure.from_container(returned_shape, recursive=True)))
computation_type = computation_types.FunctionType(parameter_type,
returned_type_spec)
return xla_serialization.create_xla_tff_computation(
compiled_xla, computation_type)
def test_set_local_execution_context(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.ComputationImpl(comp_pb, ctx_stack)
execution_contexts.set_local_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_create_xla_tff_computation(self):
xla_comp = _make_test_xla_comp()
comp_pb = xla_serialization.create_xla_tff_computation(
xla_comp, computation_types.FunctionType(None, np.int32))
self.assertIsInstance(comp_pb, pb.Computation)
self.assertEqual(comp_pb.WhichOneof('computation'), 'xla')
type_spec = type_serialization.deserialize_type(comp_pb.type)
self.assertEqual(str(type_spec), '( -> int32)')
xla_comp = xla_serialization.unpack_xla_computation(
comp_pb.xla.hlo_module)
self.assertIn('ROOT constant.1 = s32[] constant(10)',
xla_comp.as_hlo_text())
self.assertEqual(str(comp_pb.xla.parameter), '')
self.assertEqual(str(comp_pb.xla.result), 'tensor {\n'
' index: 0\n'
'}\n')
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, 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 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.get_event_loop().run_until_complete(
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.get_event_loop().run_until_complete(
ex.create_call(comp_val))
self.assertIsInstance(call_val, executor.XlaValue)
self.assertEqual(str(call_val.type_signature), 'int32')
result = asyncio.get_event_loop().run_until_complete(
call_val.compute())
self.assertEqual(result, 10)
def serialize_jax_computation(traced_fn, arg_fn, parameter_type,
context_stack):
"""Serializes a Python function containing JAX code as a TFF computation.
Args:
traced_fn: The Python function containing JAX code to be traced by JAX and
serialized as a TFF computation containing XLA code.
arg_fn: An unpacking function that takes a TFF argument, and returns a combo
of (args, kwargs) to invoke `traced_fn` with (e.g., as the one constructed
by `function_utils.create_argument_unpacking_fn`).
parameter_type: An instance of `computation_types.Type` that represents the
TFF type of the computation parameter, or `None` if the function does not
take any parameters.
context_stack: The context stack to use during serialization.
Returns:
An instance of `pb.Computation` with the constructed computation.
Raises:
TypeError: if the arguments are of the wrong types.
"""
py_typecheck.check_callable(traced_fn)
py_typecheck.check_callable(arg_fn)
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
if parameter_type is not None:
parameter_type = computation_types.to_type(parameter_type)
packed_arg = _tff_type_to_xla_serializer_arg(parameter_type)
else:
packed_arg = None
args, kwargs = arg_fn(packed_arg)
# While the fake parameters are fed via args/kwargs during serialization,
# it is possible for them to get reorderd in the actual generate XLA code.
# We use here the same flatenning function as that one, which is used by
# the JAX serializer to determine the orderding and allow it to be captured
# in the parameter binding. We do not need to do anything special for the
# results, since the results, if multiple, are always returned as a tuple.
flattened_obj, _ = jax.tree_util.tree_flatten((args, kwargs))
tensor_indexes = list(np.argsort([x.tensor_index for x in flattened_obj]))
def _adjust_arg(x):
if isinstance(x, structure.Struct):
return type_conversions.type_to_py_container(x, x.type_signature)
else:
return x
args = [_adjust_arg(x) for x in args]
kwargs = {k: _adjust_arg(v) for k, v in kwargs.items()}
context = jax_computation_context.JaxComputationContext()
with context_stack.install(context):
tracer_callable = jax.xla_computation(traced_fn,
tuple_args=True,
return_shape=True)
compiled_xla, returned_shape = tracer_callable(*args, **kwargs)
if isinstance(returned_shape, jax.ShapeDtypeStruct):
returned_type_spec = _jax_shape_dtype_struct_to_tff_tensor(
returned_shape)
else:
returned_type_spec = computation_types.to_type(
structure.map_structure(
_jax_shape_dtype_struct_to_tff_tensor,
structure.from_container(returned_shape, recursive=True)))
computation_type = computation_types.FunctionType(parameter_type,
returned_type_spec)
return xla_serialization.create_xla_tff_computation(
compiled_xla, tensor_indexes, computation_type)
